Repository: pytorch
Issue: CTCLoss cuda backend large batch handling takes up to 1.8x more memory
Labels: module: loss,module: cuda,module: memory usage,triaged,enhancement

## 🐛 Bug

The special large batch / alphabet handling, although can provide up to 2.5x speedup sometimes, it comes at the cost of up to 1.8x more memory. For large targets, this can be a significant amount of increased memory usage (in the order of multiple GBs).
I think if the memory consumption can't be put down, then this should be optional, many users with tight memory constraints would prefer a slower implementation to a memory hungry one

## To Reproduce

To make a fair comparison of the effect, I changed [this line](https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/LossCTC.cu#L642) from `if (is_large)` to `if (zero_infinity)` to make it simple to try setting `is_large` on & off without the hassle of adding a new argument, given the argument `zero_infinity` has negligible memory / time effects on other parts of the code

Then, I used the following snippet
```
from torch import nn
import torch
import math
import time

time_step  = 5  # Input sequence length
batch_size = 4  # Batch size
target_sql = 3  # Target sequence length

iters = 5
warm  = 3

ctc_loss  = [nn.CTCLoss(reduction='sum',zero_infinity=False)]
ctc_loss += [nn.CTCLoss(reduction='sum',zero_infinity=True)]

for j in range(15):
  time_step  *= 2
  target_sql *= 2
  print('Logits length : ',time_step,'Label length : ',target_sql)
  print('------------------------------------------')
  
  
  x_lengths = torch.full(size=(batch_size,), fill_value=time_step, dtype=torch.long).cuda()
  y_lengths = torch.full(size=(batch_size,), fill_value=target_sql, dtype=torch.long).cuda()

  mem   = [0,0]
  etime = [0,0]

  for r in [3,200,5000]: #vocab sizes
    print('Vocabulary : ',r)
    print('=================')
    vocab_size = r

    x = torch.randn(time_step, batch_size, vocab_size).log_softmax(2).detach().requires_grad_().cuda()
    y = torch.randint(low=1, high=vocab_size-1, size=(batch_size, target_sql),dtype=torch.long).cuda()

    for c in range(2):
      s=0
      for i in range(iters):
          torch.cuda.synchronize()
          a = time.perf_counter()
          loss = ctc_loss[c](x, y, x_lengths, y_lengths).mean()
          loss.backward()
          torch.cuda.synchronize()
          b = time.perf_counter()
          del loss
          
          if i>=warm:
              s += b - a

      mem[c] = torch.cuda.max_memory_allocated()/2**30
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.empty_cache()

      etime[c] = s/(iters-warm)
      print(f'Average [{c}] Time : {etime[c]:.5f} sec Memory : {mem[c]:.5f} GB')
  
    print(f'Speed up : {etime[0]/etime[1]:.3f} Memory increase ratio : {mem[1]/mem[0]:.3f} \n')
```
Which produced the following output.
Note that for each input / target length, we try both without (first line) and with (second line) large batch handling, then both the speed-up and increase on memory consumption is printed. This is tried for vocabularies of sizes [3,200,5000]

```
Logits length :  10 Label length :  6
------------------------------------------
Vocabulary :  3
=================
Average [0] Time : 0.00162 sec Memory : 0.00001 GB
Average [1] Time : 0.00368 sec Memory : 0.00002 GB
Speed up : 0.439 Memory increase ratio : 1.500 

Vocabulary :  200
=================
Average [0] Time : 0.00172 sec Memory : 0.00007 GB
Average [1] Time : 0.00333 sec Memory : 0.00010 GB
Speed up : 0.516 Memory increase ratio : 1.445 

Vocabulary :  5000
=================
Average [0] Time : 0.00466 sec Memory : 0.00150 GB
Average [1] Time : 0.00483 sec Memory : 0.00225 GB
Speed up : 0.966 Memory increase ratio : 1.497 

Logits length :  20 Label length :  12
------------------------------------------
Vocabulary :  3
=================
Average [0] Time : 0.00167 sec Memory : 0.00075 GB
Average [1] Time : 0.00335 sec Memory : 0.00004 GB
Speed up : 0.500 Memory increase ratio : 0.048 

Vocabulary :  200
=================
Average [0] Time : 0.00194 sec Memory : 0.00014 GB
Average [1] Time : 0.00393 sec Memory : 0.00020 GB
Speed up : 0.493 Memory increase ratio : 1.435 

Vocabulary :  5000
=================
Average [0] Time : 0.00618 sec Memory : 0.00300 GB
Average [1] Time : 0.00630 sec Memory : 0.00449 GB
Speed up : 0.981 Memory increase ratio : 1.497 

Logits length :  40 Label length :  24
------------------------------------------
Vocabulary :  3
=================
Average [0] Time : 0.00165 sec Memory : 0.00149 GB
Average [1] Time : 0.00327 sec Memory : 0.00012 GB
Speed up : 0.505 Memory increase ratio : 0.078 

Vocabulary :  200
=================
Average [0] Time : 0.00224 sec Memory : 0.00030 GB
Average [1] Time : 0.00389 sec Memory : 0.00042 GB
Speed up : 0.576 Memory increase ratio : 1.397 

Vocabulary :  5000
=================
Average [0] Time : 0.00951 sec Memory : 0.00602 GB
Average [1] Time : 0.00920 sec Memory : 0.00901 GB
Speed up : 1.033 Memory increase ratio : 1.495 

Logits length :  80 Label length :  48
------------------------------------------
Vocabulary :  3
=================
Average [0] Time : 0.00176 sec Memory : 0.00299 GB
Average [1] Time : 0.00332 sec Memory : 0.00042 GB
Speed up : 0.530 Memory increase ratio : 0.142 

Vocabulary :  200
=================
Average [0] Time : 0.00294 sec Memory : 0.00071 GB
Average [1] Time : 0.00411 sec Memory : 0.00095 GB
Speed up : 0.716 Memory increase ratio : 1.335 

Vocabulary :  5000
=================
Average [0] Time : 0.01783 sec Memory : 0.01216 GB
Average [1] Time : 0.01603 sec Memory : 0.01812 GB
Speed up : 1.112 Memory increase ratio : 1.490 

Logits length :  160 Label length :  96
------------------------------------------
Vocabulary :  3
=================
Average [0] Time : 0.00310 sec Memory : 0.00597 GB
Average [1] Time : 0.00341 sec Memory : 0.00164 GB
Speed up : 0.911 Memory increase ratio : 0.275 

Vocabulary :  200
=================
Average [0] Time : 0.00377 sec Memory : 0.00188 GB
Average [1] Time : 0.00400 sec Memory : 0.00258 GB
Speed up : 0.943 Memory increase ratio : 1.371 

Vocabulary :  5000
=================
Average [0] Time : 0.03427 sec Memory : 0.02477 GB
Average [1] Time : 0.02548 sec Memory : 0.03669 GB
Speed up : 1.345 Memory increase ratio : 1.481 

Logits length :  320 Label length :  192
------------------------------------------
Vocabulary :  3
=================
Average [0] Time : 0.00293 sec Memory : 0.01194 GB
Average [1] Time : 0.00376 sec Memory : 0.00648 GB
Speed up : 0.779 Memory increase ratio : 0.543 

Vocabulary :  200
=================
Average [0] Time : 0.00505 sec Memory : 0.00559 GB
Average [1] Time : 0.00657 sec Memory : 0.00836 GB
Speed up : 0.769 Memory increase ratio : 1.496 

Vocabulary :  5000
=================
Average [0] Time : 0.05630 sec Memory : 0.05137 GB
Average [1] Time : 0.03410 sec Memory : 0.07521 GB
Speed up : 1.651 Memory increase ratio : 1.464 

Logits length :  640 Label length :  384
------------------------------------------
Vocabulary :  3
=================
Average [0] Time : 0.00402 sec Memory : 0.02388 GB
Average [1] Time : 0.00386 sec Memory : 0.02577 GB
Speed up : 1.042 Memory increase ratio : 1.079 

Vocabulary :  200
=================
Average [0] Time : 0.00852 sec Memory : 0.01850 GB
Average [1] Time : 0.00817 sec Memory : 0.02953 GB
Speed up : 1.042 Memory increase ratio : 1.596 

Vocabulary :  5000
=================
Average [0] Time : 0.09079 sec Memory : 0.11005 GB
Average [1] Time : 0.08103 sec Memory : 0.15774 GB
Speed up : 1.121 Memory increase ratio : 1.433 

Logits length :  1280 Label length :  768
------------------------------------------
Vocabulary :  3
=================
Average [0] Time : 0.01199 sec Memory : 0.05877 GB
Average [1] Time : 0.00849 sec Memory : 0.10569 GB
Speed up : 1.412 Memory increase ratio : 1.798 

Vocabulary :  200
=================
Average [0] Time : 0.02243 sec Memory : 0.06629 GB
Average [1] Time : 0.01753 sec Memory : 0.11320 GB
Speed up : 1.279 Memory increase ratio : 1.708 

Vocabulary :  5000
=================
Average [0] Time : 0.22960 sec Memory : 0.25007 GB
Average [1] Time : 0.16240 sec Memory : 0.34577 GB
Speed up : 1.414 Memory increase ratio : 1.383 

Logits length :  2560 Label length :  1536
------------------------------------------
Vocabulary :  3
=================
Average [0] Time : 0.03153 sec Memory : 0.23473 GB
Average [1] Time : 0.02524 sec Memory : 0.41070 GB
Speed up : 1.249 Memory increase ratio : 1.750 

Vocabulary :  200
=================
Average [0] Time : 0.08150 sec Memory : 0.24976 GB
Average [1] Time : 0.04038 sec Memory : 0.42573 GB
Speed up : 2.018 Memory increase ratio : 1.705 

Vocabulary :  5000
=================
Average [0] Time : 0.39866 sec Memory : 0.61731 GB
Average [1] Time : 0.39342 sec Memory : 0.80872 GB
Speed up : 1.013 Memory increase ratio : 1.310 

Logits length :  5120 Label length :  3072
------------------------------------------
Vocabulary :  3
=================
Average [0] Time : 0.19747 sec Memory : 0.93821 GB
Average [1] Time : 0.08911 sec Memory : 1.64171 GB
Speed up : 2.216 Memory increase ratio : 1.750 

Vocabulary :  200
=================
Average [0] Time : 0.35623 sec Memory : 0.96900 GB
Average [1] Time : 0.13945 sec Memory : 1.67250 GB
Speed up : 2.554 Memory increase ratio : 1.726 

Vocabulary :  5000
=================
Average [0] Time : 1.09736 sec Memory : 1.70069 GB
Average [1] Time : 0.93109 sec Memory : 2.40419 GB
Speed up : 1.179 Memory increase ratio : 1.414 

Logits length :  10240 Label length :  6144
------------------------------------------
Vocabulary :  3
=================
Average [0] Time : 1.11494 sec Memory : 3.75141 GB
Average [1] Time : 0.65411 sec Memory : 6.56467 GB
Speed up : 1.705 Memory increase ratio : 1.750 

Vocabulary :  200
=================
Average [0] Time : 1.77882 sec Memory : 3.81299 GB
Average [1] Time : 0.93758 sec Memory : 6.62626 GB
Speed up : 1.897 Memory increase ratio : 1.738 

Vocabulary :  5000
=================
Average [0] Time : 3.43761 sec Memory : 5.27784 GB
Average [1] Time : 2.56783 sec Memory : 8.09110 GB
Speed up : 1.339 Memory increase ratio : 1.533 

Logits length :  20480 Label length :  12288
------------------------------------------
Vocabulary :  3
=================
RuntimeError: CUDA out of memory. Tried to allocate 7.50 GiB (GPU 0; 11.91 GiB total capacity; 7.50 GiB already allocated; 3.22 GiB free; 2.37 MiB cached)
```

## Environment

PyTorch Version (e.g., 1.0): tested problem on master
OS (e.g., Linux): Linux
How you installed PyTorch (conda, pip, source): pip
Python version: 3.6
CUDA/cuDNN version: CUDA: 10.0.130, cuDNN: 7.4.2
GPU models and configuration: tested on Titan XP

@t-vi