RAM leakage when trying to retrieve the hidden states from the GPT-2 model.
Author: h56cho
Pull request: no

Hello,

I am trying to retrieve hidden state vectors from my trained GPT-2 model from a loop, and there is a huge RAM leakage associated with the operation. Below are my code:

```python
# for loop to calculate TVD
def TVD_loop(test_i, test_dummy_i, nlayer, best_model):
    
    TVD_tensor = torch.zeros(test_i.size()[1], (nlayer+1), test_i.size()[0]).float()
    
    # replace every 0's in TVD_tensor to -2
    TVD_tensor = torch.where(TVD_tensor == 0.0, torch.tensor(-2.0), TVD_tensor)
    
    for m in range(test_i.size()[1]):

        input_ids = test_i[:,m]
        input_ids = torch.tensor(input_ids.tolist()).unsqueeze(0) 
        
        # NOTE: Hidden states are in torch.FloatTensor,
        #       (one for the output of each layer + the output of the embeddings)
        # jth layer
        for j in range(nlayer+1):
            
            tst_hidden_states = best_model(input_ids)[3][j][0, (test_i.size()[0] - 1), :]
            
            for l in range(m * test_i.size()[0], (m+1) * test_i.size()[0]):
                
                input_ids_dummy = test_dummy_i[:,l]
                input_ids_dummy = torch.tensor(input_ids_dummy.tolist()).unsqueeze(0) 
                
                tst_hidden_states_dummy = best_model(input_ids_dummy)[3][j][0, (test_i.size()[0] - 1), :]
                
                # TVD_tensor[i,j,k] denotes for TVDC calculated at 
                # batch i, layer j, and dummy output k
                TVD_tensor[m,j,(l % (test_i.size()[0]))] = TVD(tst_hidden_states, tst_hidden_states_dummy)

    return TVD_tensor 
```

I have about ~400GB of RAM, but each time the hidden state vector is retrieved (e.g. ```tst_hidden_states = best_model(input_ids)[3][j][0, (test_i.size()[0] - 1), :]```), it uses up about ~2GB of RAM. How can I prevent this? Thank you,


Comments:
To add more, I fixed my code like below:

```python
# for loop to calculate TVD
def TVD_loop(test_i, test_dummy_i, nlayer, best_model):
    
    TVD_tensor = torch.zeros(test_i.size()[1], (nlayer+1), test_i.size()[0]).float()
    
    # replace every 0's in TVD_tensor to -2
    TVD_tensor = torch.where(TVD_tensor == 0.0, torch.tensor(-2.0), TVD_tensor)
    
    for m in range(test_i.size()[1]):
        
        gc.collect()
        
        input_ids = test_i[:,m]
        input_ids = torch.tensor(input_ids.tolist()).unsqueeze(0) 

        # NOTE: Hidden states are in torch.FloatTensor,
        #       (one for the output of each layer + the output of the embeddings)
        # jth layer
        for j in range(nlayer+1):
            
            del gc.garbage[:]
            gc.collect()
            
            for l in range(m * test_i.size()[0], (m+1) * test_i.size()[0]):
                
                del gc.garbage[:]
                gc.collect()
                
                tst_hidden_states = best_model(input_ids)[3][j][0, (test_i.size()[0] - 1), :]
                
                input_ids_dummy = test_dummy_i[:,l]
                input_ids_dummy = torch.tensor(input_ids_dummy.tolist()).unsqueeze(0) 
                
                tst_hidden_states_dummy = best_model(input_ids_dummy)[3][j][0, (test_i.size()[0] - 1), :]
                
                del input_ids_dummy
                del gc.garbage[:]
                gc.collect()
                
                # TVD_tensor[i,j,k] denotes for TVD calculated at 
                # batch i, layer j, and dummy output k
                TVD_tensor[m,j,(l % (test_i.size()[0]))] = TVD(tst_hidden_states, tst_hidden_states_dummy)
                
                del tst_hidden_states
                del tst_hidden_states_dummy
                del gc.garbage[:]
                gc.collect()
            
                print('l={}, gc_get_count={}'.format(l,gc.get_count()))
                
            del gc.garbage[:]
            gc.collect()
            print('j={}, gc_get_count={}'.format(j,gc.get_count()))
            
        del gc.garbage[:]
        del input_ids
        gc.collect()
        
        print('m={}, gc_get_count={}'.format(m,gc.get_count()))
             
    return TVD_tensor      
```

from the code above, when ```m=0, j=0, l=0```, everything is fine, but once ```m=0, j=1, l=0``` is reached, the memory usage starts to accumulate rapidly. How should I fix my code?

Did you try detaching the gradient from the hidden states by replacing 
`
tst_hidden_states = best_model(input_ids)[3][j][0, (test_i.size()[0] - 1), :]
`
with
`
tst_hidden_states = best_model(input_ids)[3][j][0, (test_i.size()[0] - 1), :].detach()
`
?
The gradient needed for backpropagation usually consumes a lot of RAM and is probably not needed in your case.

> 

> Did you try detaching the gradient from the hidden states by replacing
> `tst_hidden_states = best_model(input_ids)[3][j][0, (test_i.size()[0] - 1), :]`
> with
> `tst_hidden_states = best_model(input_ids)[3][j][0, (test_i.size()[0] - 1), :].detach()`
> ?
> The gradient needed for backpropagation usually consumes a lot of RAM and is probably not needed in your case.
> 
> >

Detaching _before_ the slice might be even better? Not sure.

Thank you! .detach() solved this RAM leakage issue :)