ASR Task Failing due to CUDA Memory Issue - How to introduce Lightning Fabric support for Lightning Flash Tasks?

[greeshmasmenon]

Hi,

I am trying to finetune the Wav2vec2 model ("facebook/wav2vec2-large-960h-lv60-self") with custom data that i have.

GPU : Tesla V100-SXM2-16GB
Number of GPUs: 8 (2 Nodes of 4 each)
Shape of audio dataset (Each audio segment is roughly 3 seconds long) : { "training": [ 133328, 4 ], "validation": [ 33332, 3 ] }

Each audio segment is roughly 3 seconds long. The training arguments are below :

class TrainingArgs:
    TRAIN_FILE_PATH: str = "/home/users/gmenon/notebooks/home/users/gmenon/notebooks/train_metadata_cleaned.csv"
    TEST_FILE_PATH: str = "/home/users/gmenon/notebooks/home/users/gmenon/notebooks/validation_metadata_cleaned.csv"
    MODEL_BACKBONE: str = "facebook/wav2vec2-large-960h-lv60-self"
    BATCH_SIZE: int = 1
    NUM_EPOCHS: int = 20
    NUM_GPUS = torch.cuda.device_count()
    MODEL_SAVE_PATH: str = "model_artefacts/finetuned_ASR_model.pt"
    FINETUNE_STRATEGY: str = "freeze"
    ACCUMULATE_GRAD_BATCHES = 1
    PRECISION = 16
    MAX_EPOCHS = 10 

I am getting the CUDA Memory error -

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 250.00 MiB (GPU 0; 15.77 GiB total capacity; 14.23 GiB already allocated; 183.38 MiB free; 14.43 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

I would like to use a distributed training approach as I have gone down to BATCH SIZE of 1 and ACCUMULATE_GRAD_BATCHES = 1 and don't know how to reduce the data loaded to the GPU still further. Any advice here would be appreciated?

Also, I want to start looking at Lightning Fabric and introduce the parallel training procedures to see if it solves my problem. With the high level interfaces, I am not sure where to start. Can someone guide me how?

Sun Jul 30 00:39:33 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.85.12    Driver Version: 525.85.12    CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla V100-SXM2...  On   | 00000000:1A:00.0 Off |                    0 |
| N/A   45C    P0    43W / 300W |      0MiB / 16384MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Tesla V100-SXM2...  On   | 00000000:1C:00.0 Off |                    0 |
| N/A   40C    P0    43W / 300W |      0MiB / 16384MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   2  Tesla V100-SXM2...  On   | 00000000:1D:00.0 Off |                    0 |
| N/A   38C    P0    42W / 300W |      0MiB / 16384MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   3  Tesla V100-SXM2...  On   | 00000000:1E:00.0 Off |                    0 |
| N/A   40C    P0    43W / 300W |      0MiB / 16384MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

Logs:

Epoch 0:   8%|▊         | 1254/16666 [02:12<27:03,  9.49it/s, loss=892, v_num=3, train_loss_step=916.0]
Epoch 0:   8%|▊         | 1255/16666 [02:12<27:03,  9.49it/s, loss=892, v_num=3, train_loss_step=916.0]
Epoch 0:   8%|▊         | 1255/16666 [02:12<27:03,  9.49it/s, loss=877, v_num=3, train_loss_step=496.0]
Epoch 0:   8%|▊         | 1256/16666 [02:12<27:03,  9.49it/s, loss=877, v_Traceback (most recent call last):
  File "/home/users/gmenon/workspace/songsLyricsGenerator/src/wav2vec2_training.py", line 9, in <module>
    wav2vec2.finetune()
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/src/training/wav2vec2_finetune.py", line 63, in finetune
    self.wav2vec2_trainer.finetune(self.wav2vec2_model,
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/flash/core/trainer.py", line 164, in finetune
    return super().fit(model, train_dataloader, val_dataloaders, datamodule)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 608, in fit
    call._call_and_handle_interrupt(
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/trainer/call.py", line 38, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 650, in _fit_impl
    self._run(model, ckpt_path=self.ckpt_path)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1112, in _run
    results = self._run_stage()
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1191, in _run_stage
    self._run_train()
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1214, in _run_train
    self.fit_loop.run()
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/loops/fit_loop.py", line 267, in advance
    self._outputs = self.epoch_loop.run(self._data_fetcher)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 213, in advance
    batch_output = self.batch_loop.run(kwargs)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py", line 88, in advance
    outputs = self.optimizer_loop.run(optimizers, kwargs)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 202, in advance
    result = self._run_optimization(kwargs, self._optimizers[self.optim_progress.optimizer_position])
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 249, in _run_optimization
    self._optimizer_step(optimizer, opt_idx, kwargs.get("batch_idx", 0), closure)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 370, in _optimizer_step
    self.trainer._call_lightning_module_hook(
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1356, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/core/module.py", line 1754, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py", line 169, in step
    step_output = self._strategy.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/strategies/ddp.py", line 280, in optimizer_step
    optimizer_output = super().optimizer_step(optimizer, opt_idx, closure, model, **kwargs)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/strategies/strategy.py", line 234, in optimizer_step
    return self.precision_plugin.optimizer_step(
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/plugins/precision/native_amp.py", line 75, in optimizer_step
    closure_result = closure()
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 149, in __call__
    self._result = self.closure(*args, **kwargs)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 135, in closure
    step_output = self._step_fn()
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 419, in _training_step
    training_step_output = self.trainer._call_strategy_hook("training_step", *kwargs.values())
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1494, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/strategies/ddp.py", line 351, in training_step
    return self.model(*args, **kwargs)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 1156, in forward
    output = self._run_ddp_forward(*inputs, **kwargs)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 1110, in _run_ddp_forward
    return module_to_run(*inputs[0], **kwargs[0])  # type: ignore[index]
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/pytorch_lightning/overrides/base.py", line 98, in forward
    output = self._forward_module.training_step(*inputs, **kwargs)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/flash/core/model.py", line 437, in training_step
    output = self.step(batch, batch_idx, self.train_metrics)
  File "/mnt/irisgpfs/users/gmenon/workspace/songsLyricsGenerator/slg_wav2vec2/lib/python3.8/site-packages/flash/audio/speech_recognition/model.py", line 97, in step
    out = self.model(batch["input_values"], labels=batch["labels"])