CUDA out of memory while training Entity Linking

[rdemorais]

I'm trying to train an Entity Linking model using the following configuration:

[paths]
train = "corpus/shc_el_train.spacy"
dev = "corpus/shc_el_dev.spacy"
kb = "assets/shc_el_kb"

[system]
seed = 342
gpu_allocator = "pytorch"

[nlp]
lang = "en"
pipeline = ["transformer", "sentencizer", "entity_linker"]
batch_size = 8
disabled = []
before_creation = null
after_creation = null
after_pipeline_creation = null
tokenizer = {"@tokenizers":"spacy.Tokenizer.v1"}

[components]

[components.transformer]
factory = "transformer"
set_extra_annotations = {"@annotation_setters":"spacy-transformers.null_annotation_setter.v1"}

[components.transformer.model]
@architectures = "spacy-transformers.TransformerModel.v3"
name = "assets/shc-cn-v2"
tokenizer_config = {"use_fast": true}

[components.sentencizer]
factory = "sentencizer"
punct_chars = null

[components.entity_linker]
factory = "entity_linker"
entity_vector_length = 768
get_candidates = {"@misc":"spacy.CandidateGenerator.v1"}
incl_context = true
incl_prior = false
labels_discard = ["DISP", "ESC", "ESPE", "ESTA", "EXRES", "NOME", "SINV", "SUV"]
n_sents = 1
overwrite = true
scorer = {"@scorers":"spacy.entity_linker_scorer.v1"}

[components.entity_linker.model]
@architectures = "spacy.EntityLinker.v2"
nO = null

[components.entity_linker.model.tok2vec]
@architectures = "spacy-transformers.TransformerListener.v1"
grad_factor = 1
upstream = "*"

[components.entity_linker.model.tok2vec.pooling]
@layers = "reduce_mean.v1"

[initialize]
vectors = ${paths.vectors}

[initialize.components]

[initialize.components.entity_linker]

[initialize.components.entity_linker.kb_loader]
@misc = "spacy.KBFromFile.v1"
kb_path = ${paths.kb}

[corpora]

[corpora.train]
@readers = "spacy.Corpus.v1"
file = ${paths.train}

[corpora.dev]
@readers = "spacy.Corpus.v1"
file = ${paths.dev}

[training]
train_corpus = "corpora.train"
dev_corpus = "corpora.dev"
seed = ${system.seed}
gpu_allocator = ${system.gpu_allocator}
dropout = 0.2
patience = 1000
eval_frequency = 200
accumulate_gradient = 2
max_epochs = 0
max_steps = 600
frozen_components = ["sentencizer"]
annotating_components = []
before_to_disk = null

[training.batcher]
@batchers = "spacy.batch_by_words.v1"
discard_oversize = false
tolerance = 0.2
get_length = null

[training.batcher.size]
@schedules = "compounding.v1"
start = 4
stop = 32
compound = 1.001
t = 0.0

[training.optimizer]
@optimizers = "Adam.v1"
beta1 = 0.9
beta2 = 0.999
L2_is_weight_decay = true
L2 = 0.01
grad_clip = 1.0
use_averages = false
eps = 0.00000001
learn_rate = 0.001

[training.score_weights]
nel_micro_p = 0.0
nel_micro_r = 0.0
nel_micro_f = 1.0

[pretraining]

[optimizer]
@optimizers = "Adam.v1"
learn_rate = 0.001
beta1 = 0.9
beta2 = 0.999
L2 = 0.0
eps = 0.00000001
grad_clip = 1.0
L2_is_weight_decay = true
use_averages = true

The transformer used can be replaced to anyone with MASK task training. The knowledge base has 26k entries, the train/dev sets (527 docs / 127 docs).

I'm able to train the Entity Linking using the script provided by the NEL Emerson example, by when I switch to the config approach, the error occurs.

What I have already tried:

1. Change the [nlp] batch_size to 32, 16, 8...
2. Change the values of [training.batcher.size].

How to reproduce the behaviour

Execute python -m spacy train ${vars.trf_config_final} --output training/ --gpu-id 0

Your Environment

• spaCy version: 3.3.0
• Platform: Linux-5.4.188+-x86_64-with-Ubuntu-18.04-bionic
• Python version: 3.7.13
• Pipelines: en_core_web_sm (3.3.0)
• Tesla P100-PCIE-16GB

ERROR Description:

RuntimeError: CUDA out of memory. Tried to allocate 1.99 GiB (GPU 0; 15.90 GiB total capacity; 13.73 GiB already allocated; 829.75 MiB free; 14.19 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

[rdemorais]

UPDATE

I've changed the above [training.batcher] discard_oversize = false to true. Since I'm using Colab, this time I've got a Tesla T4 GPU and I could resume the training. Not sure if it is due to the parameter or the GPU.

[polm]

I've changed the above [training.batcher] discard_oversize = false to true.

Note this means you will discard items larger than the batch size, which is 100 tokens by default, so that's probably not what you want.

How long are your documents? If they're very long you may need to split them up to keep memory usage down. It may also be an issue that your KB is taking up GPU memory and leaves you with not enough left to train effectively. If that's the case you'll have to reduce your KB size or get more memory. Have you considered running a CPU pipeline instead of Transformers to confirm that works?

[polm]

Also, please don't post screenshots of code or terminal output, post them as text.

[polm]

For reference, note this is a continuation of part of #7892.