CUDA out of memory when transcribing long audio
State: closed | Pull request: no

I'm trying to transcribe some audio files using NeMo/examples/asr/transcribe_speech_parallel.py using the stt_en_conformer_transducer_small model with these parameters:
`predict_ds.batch_size=1 
  trainer.precision=16 
  trainer.devices=1 
  predict_ds.num_workers=1`

I'm using just one GPU for debug purposes. It's an NVIDIA RTX A6000. I'm able to transcribe audio files up to a couple hundred seconds long, but I'm getting the following error once the decoding reaches a file which is 1025 seconds long:

 

> Error executing job with overrides: ['model=stt_en_conformer_transducer_small', 'predict_ds.manifest_filepath=dataset.json', 
> 'predict_ds.batch_size=1', 'trainer.precision=16', 'trainer.devices=1', 'predict_ds.num_workers=1', 'output_path=output_dir']
> Traceback (most recent call last):
>   File "NeMo/examples/asr/transcribe_speech_parallel.py", line 153, in main
>     predictions = trainer.predict(model=model, dataloaders=data_loader, return_predictions=cfg.return_predictions)
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 949, in predict
>     return self._call_and_handle_interrupt(
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 648, in _call_and_handle_interrupt
>     return self.strategy.launcher.launch(trainer_fn, *args, trainer=self, **kwargs)
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 93, in launch
>     return function(*args, **kwargs)
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 996, in _predict_impl
>     results = self._run(model, ckpt_path=self.ckpt_path)
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1166, in _run
>     results = self._run_stage()
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1251, in _run_stage
>     return self._run_predict()
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1311, in _run_predict
>     return self.predict_loop.run()
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 200, in run
>     self.advance(*args, **kwargs)
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/loops/dataloader/prediction_loop.py", line 102, in advance
>     dl_predictions, dl_batch_indices = self.epoch_loop.run(
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 200, in run
>     self.advance(*args, **kwargs)
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/prediction_epoch_loop.py", line 101, in advance
>     self._predict_step(batch, batch_idx, dataloader_idx)
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/prediction_epoch_loop.py", line 130, in _predict_step
>     predictions = self.trainer._call_strategy_hook("predict_step", *step_kwargs.values())
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1704, in _call_strategy_hook
>     output = fn(*args, **kwargs)
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/strategies/ddp.py", line 369, in predict_step
>     return self.model.predict_step(*args, **kwargs)
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/nemo/collections/asr/models/rnnt_models.py", line 767, in predict_step
>     encoded, encoded_len = self.forward(input_signal=signal, input_signal_length=signal_len)
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/nemo/core/classes/common.py", line 1084, in __call__
>     outputs = wrapped(*args, **kwargs)
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/nemo/collections/asr/models/rnnt_models.py", line 663, in forward
>     encoded, encoded_len = self.encoder(audio_signal=processed_signal, length=processed_signal_length)
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
>     return forward_call(*input, **kwargs)
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/nemo/core/classes/common.py", line 1084, in __call__
>     outputs = wrapped(*args, **kwargs)
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/nemo/collections/asr/modules/conformer_encoder.py", line 424, in forward
>     audio_signal = layer(
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
>     return forward_call(*input, **kwargs)
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/nemo/collections/asr/parts/submodules/conformer_modules.py", line 140, in forward
>     x = self.self_attn(
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
>     return forward_call(*input, **kwargs)
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/nemo/collections/asr/parts/submodules/multi_head_attention.py", line 244, in forward
>     out = self.forward_attention(v, scores, mask)
>   File "/home/msancinetti/miniconda3/envs/nemo/lib/python3.8/site-packages/nemo/collections/asr/parts/submodules/multi_head_attention.py", line 109, in forward_attention
>     attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0)  # (batch, head, time1, time2)
> RuntimeError: CUDA out of memory. Tried to allocate 10.00 GiB (GPU 0; 47.54 GiB total capacity; 36.63 GiB already allocated; 8.31 GiB free; 37.49 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
> 

My guess is that the model is doing attention on the whole audio and that blows up in memory which is why I'm only having this issue for longer audio files but not for shorter ones. Is there a way to limit the maximum attention context?