Repository: pytorch
Issue: Memory inefficient in batched matmul when requiring gradients
Labels: triaged,module: linear algebra

While implementing the batched matrix multiplication, i noticed that the batched matrix multiplication is not efficient, see the code below

```
import torch                                                                                                                                                                             
                                                                                                                                                                                         
# Input tensor                                                                                                                                                                           
## Batch size=8192, dim=512                                                                                                                                                             
x = torch.FloatTensor(8192, 512).requires_grad_().cuda()                                                                                                                                 
                                                                                                                                                                                         
if True:                                                                                                                                                                                 
    # Batch strategy 1                                                                                                                                                                   
    x1 = x.view(8192, 8, 1, 64)  # 512 = 8 * 64                                                                                                                                        
    W1 = torch.FloatTensor(8, 64, 64).cuda()                                                                                                                                             
    out1 = torch.matmul(x1, W1)  # out: [8192, 8, 1, 64]                                                                                                                                
                                                                                                                                                                                         
    print(torch.cuda.memory_allocated())    # 1107427328                                                                                                                                                
                                                                                                                                                                                         
if False:                                                                                                                                                                                
    # Batch strategy 2                                                                                                                                                                   
    x2 = x.view(8192, 1, 512)  # add one dimension for batch matmul                                                                                                                      
    W2 = torch.FloatTensor(512, 512).cuda()  # larger than W1                                                                                                                            
    # out: [8192, 1, 1024] # the same number of elements as out1                                                                                                                                   
    out2 = torch.matmul(x2, W2)                                                                                                                                                          
    print(torch.cuda.memory_allocated())  # 34603008
```
However, it turns out that Batch strategy 2 has less memory cost despite that W2 is larger than Batch strategy 1. And everything else are the same (x1, x2 have same number of elements, also out1, out2).

I also found that by removing the requires_grad_() the memory costs are similar (~33685504).

What’s the possible reason for that? (Not sure if this is a real issue, so I also posted on the forum).

cc @jianyuh @nikitaved @pearu @mruberry @heitorschueroff @walterddr @IvanYashchuk @xwang233 @Lezcano