Error at loss.backward() when pretraining Llama model from scratch "TRYING TO BACKWARD SECOND TIME"

abbas_haider: <p>I m trying to explore Llama code at <a href="https://github.com/facebookresearch/llama" rel="noopener nofollow ugc">https://github.com/facebookresearch/llama</a>.</p>
<p>My training script is shown below:</p>
<p>from datasets import load_dataset<br>
from fairscale.nn.model_parallel.initialize import (<br>
initialize_model_parallel<br>
)<br>
from llama.model_modified import Transformer_modified, ModelArgs<br>
from torch.utils.data import Dataset, DataLoader<br>
from transformers import AutoTokenizer<br>
import os<br>
import torch<br>
import torch.distributed as dist<br>
from torch.nn.parallel import DistributedDataParallel as DDP<br>
import torch.nn.functional as F<br>
import pdb</p>
<p>class Model:<br>
def <strong>init</strong>(self, model, tokenizer):<br>
self.model = model<br>
self.tokenizer = tokenizer<br>
self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=0.001)</p>
<pre><code>def custom_loss_function(self, outputs, labels):
    softmax_outputs = F.softmax(outputs, dim=-1)

    # Flatten both the softmax_outputs and labels
    softmax_outputs_flat = softmax_outputs.view(-1, softmax_outputs.size(-1))
    labels_flat = labels.view(-1)

    # Compute the cross-entropy loss
    loss = F.cross_entropy(softmax_outputs_flat, labels_flat, reduction='mean')
    return loss

def mask_tokens(self, input_tensor, mask_prob=0.15):
    """
    Randomly masks tokens in the input tensor for masked language modeling pretraining.

    Args:
        input_tensor (torch.Tensor): Tensor containing input token IDs.
        tokenizer: Pretrained tokenizer.
        mask_prob (float): Probability of masking a token.

    Returns:
        (torch.Tensor, torch.Tensor): Masked input tensor, labels tensor.
    """
    mask = torch.rand(input_tensor.shape) &lt; mask_prob
    masked_tensor = input_tensor.clone()
    masked_tensor[mask] = self.tokenizer.mask_token_id

    # Prepare labels tensor for computing loss
    labels = torch.full_like(input_tensor, fill_value=-100)  # -100 is the default value for ignored index in cross-entropy loss
    labels[mask] = input_tensor[mask]

    return masked_tensor, labels

def count_parameters(self):
    count = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
    count = count / 1000**2
    return count

def train_model(self, data_loader):
    # Perform model training here
    for index, (tensor_input, text) in enumerate(data_loader):
        print(f'Index : {index+1}')
        # print(f'Input : {text}')
        self.optimizer.zero_grad()
        tensor_input, targets = self.mask_tokens(tensor_input.squeeze(0))
        # # Forward pass
        print(f'Input : {tensor_input} has grad : {tensor_input.requires_grad}')
        output = self.model(tokens = tensor_input, start_pos = 0)
        print(f'Output : {output} has grad : {output.requires_grad}')
        # # Loss computation
        loss = self.custom_loss_function(output, targets)
        print(f'Loss : {loss.item()}')
        print(f'Loss has grad : {loss.requires_grad}')
        # print(output._version)
        # ## Backpropagation
        loss.backward()
        # # Weight updates
        self.optimizer.step()
</code></pre>
<p>class CustomDataset(Dataset):<br>
def <strong>init</strong>(self, data, tokenizer, max_length=512):<br>
self.data = data<br>
self.tokenizer = tokenizer<br>
self.sequence_length = max_length</p>
<pre><code>def __len__(self):
    return len(self.data)

def __getitem__(self, index):
    # Get data at the specified index
    self.input_ids = self.tokenizer(self.data['train'][index]['text'], return_tensors="pt", max_length=self.sequence_length, truncation=True, padding='max_length')['input_ids']
    return self.input_ids, self.data['train'][index]['text']
</code></pre>
<p>def main():<br>
# Load datasets from the datasets library<br>
tokenizer = AutoTokenizer.from_pretrained(‘bert-base-uncased’)<br>
data = load_dataset(‘rotten_tomatoes’)<br>
sequence_length = 32<br>
dataset = CustomDataset(data=data, tokenizer=tokenizer, max_length=sequence_length)<br>
data_loader = DataLoader(dataset, batch_size=1, shuffle=True)</p>
<pre><code># Initialize model parallel
if not torch.distributed.is_initialized():
    torch.distributed.init_process_group(backend="nccl")
initialize_model_parallel(1)
torch.autograd.set_detect_anomaly(True)
model_args = ModelArgs(dim=32, vocab_size=tokenizer.vocab_size, n_layers=1, n_heads=1, max_seq_len=sequence_length,)
# pdb.set_trace()
m = Model(model=Transformer_modified(model_args), tokenizer=tokenizer)
print(f'Model has {m.count_parameters():.1f}M trainable parameters.')
# model = DDP(model, device_ids=[0], output_device=0)
m.train_model(data_loader)
# clean up
dist.destroy_process_group()
</code></pre>
<p>if <strong>name</strong> == “<strong>main</strong>”:<br>
main()</p>
<p>I got the RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.</p>
<p>I can’t find the issue. I have removed the decorator “<span class="mention">@torch.inference_mode</span>()” from the Transformer class inside Llama model.py for training.<br>
Can someone explain me what’s going wrong in this?</p>
<p>OUTPUT:</p>
<p>Model has 2.0M trainable parameters.<br>
Index : 1<br>
Input : tensor([[  101,  1996,  9882,  2135,  9603, 13633,  1997,  1000,  1996,  2935,<br>
1997,  1996,  7635,  1000, 11544,  2003,  2061,  4121,  2008,  1037,<br>
103,  1997,  2616,  3685, 23613,  6235,  2522,  1011,  3213,  1013,<br>
2472,   103]]) has grad : False<br>
Output : tensor([[[-0.0419, -0.7551, -0.4873,  …, -0.2259,  0.3079,  0.3085],<br>
[ 0.1105, -0.3129, -1.0512,  …,  0.7456, -0.4197,  1.0040],<br>
[-0.1352,  0.1497, -0.6379,  …, -0.0768,  1.3803,  0.8941],<br>
…,<br>
[-0.1277, -0.4814,  0.0347,  …,  0.3797, -0.1792,  1.4547],<br>
[ 0.2541,  0.9202, -0.2216,  …, -1.2154, -1.1796, -1.2808],<br>
[ 0.6918, -0.6579, -0.8302,  …,  0.5668, -0.5223,  1.2523]]],<br>
grad_fn=) has grad : True<br>
Loss : 10.326194763183594<br>
Loss has grad : True<br>
Index : 2<br>
Input : tensor([[  101,  4621,  2021,  2205,  1011,  8915, 23267, 16012, 24330,   102,<br>
0,     0,     0,     0,     0,   103,     0,     0,     0,     0,<br>
0,   103,     0,     0,     0,     0,   103,     0,     0,     0,<br>
0,     0]]) has grad : False<br>
Output : tensor([[[-0.0462, -0.7632, -0.4925,  …, -0.2175,  0.3111,  0.3207],<br>
[-0.5318,  0.5070,  0.2948,  …, -0.5086, -0.8850, -1.0930],<br>
[-0.9778, -0.4862,  0.1357,  …, -0.2190,  0.3070, -0.3706],<br>
…,<br>
[ 1.3116,  0.1696, -0.0291,  …,  0.4612, -0.2122,  0.9665],<br>
[ 1.3152,  0.1716, -0.0287,  …,  0.4622, -0.2161,  0.9684],<br>
[ 1.3191,  0.1746, -0.0282,  …,  0.4631, -0.2187,  0.9687]]],<br>
grad_fn=) has grad : True<br>
Loss : 10.326163291931152<br>
Loss has grad : True<br>
/home/3057693@eeecs.qub.ac.uk/Documents/llama/env/lib/python3.10/site-packages/torch/autograd/<strong>init</strong>.py:251: UserWarning: Error detected in torch::autograd::CopySlices. Traceback of forward call that caused the error:<br>
File “/home/3057693@eeecs.qub.ac.uk/Documents/llama/train_clean.py”, line 118, in <br>
main()<br>
File “/home/3057693@eeecs.qub.ac.uk/Documents/llama/train_clean.py”, line 113, in main<br>
m.train_model(data_loader)<br>
File “/home/3057693@eeecs.qub.ac.uk/Documents/llama/train_clean.py”, line 69, in train_model<br>
output = self.model(tokens = tensor_input, start_pos = 0)<br>
File “/home/3057693@eeecs.qub.ac.uk/Documents/llama/env/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 1518, in _wrapped_call_impl<br>
return self._call_impl(*args, **kwargs)<br>
File “/home/3057693@eeecs.qub.ac.uk/Documents/llama/env/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 1527, in _call_impl<br>
return forward_call(*args, **kwargs)<br>
File “/home/3057693@eeecs.qub.ac.uk/Documents/llama/llama/model_modified.py”, line 456, in forward<br>
h = layer(h, start_pos, freqs_cis, mask)<br>
File “/home/3057693@eeecs.qub.ac.uk/Documents/llama/env/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 1518, in _wrapped_call_impl<br>
return self._call_impl(*args, **kwargs)<br>
File “/home/3057693@eeecs.qub.ac.uk/Documents/llama/env/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 1527, in _call_impl<br>
return forward_call(*args, **kwargs)<br>
File “/home/3057693@eeecs.qub.ac.uk/Documents/llama/llama/model_modified.py”, line 372, in forward<br>
h = x + self.attention.forward(<br>
File “/home/3057693@eeecs.qub.ac.uk/Documents/llama/llama/model_modified.py”, line 258, in forward<br>
self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv<br>
(Triggered internally at …/torch/csrc/autograd/python_anomaly_mode.cpp:114.)<br>
Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass<br>
Traceback (most recent call last):<br>
File “/home/3057693@eeecs.qub.ac.uk/Documents/llama/train_clean.py”, line 118, in <br>
main()<br>
File “/home/3057693@eeecs.qub.ac.uk/Documents/llama/train_clean.py”, line 113, in main<br>
m.train_model(data_loader)<br>
File “/home/3057693@eeecs.qub.ac.uk/Documents/llama/train_clean.py”, line 77, in train_model<br>
loss.backward()<br>
File “/home/3057693@eeecs.qub.ac.uk/Documents/llama/env/lib/python3.10/site-packages/torch/_tensor.py”, line 492, in backward<br>
torch.autograd.backward(<br>
File “/home/3057693@eeecs.qub.ac.uk/Documents/llama/env/lib/python3.10/site-packages/torch/autograd/<strong>init</strong>.py”, line 251, in backward<br>
Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass<br>
RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.</p>

Why11699: <p>Hi abbash! I’m trying to modify Llama source code from github for my own project, and this error also happened to me. The first time in the loop loss.backward() works fine but error occurred the second time. Did you solve it? I hope we can get help from each other. From discussions of this error I can approach, It seems that some constant tensors may be involved in the computation of loss, which are freed from computation graph in the first time loss.backward() was called. But I can’t find any tensor like that in my code for days so I guess it’s something in Llama code. I also tried to use detach() on some variables but it didn’t work for me.</p>

abbas_haider: <p>Hello Wanghy,</p>
<p>The only way is to use the “retain_graph=True” flag alongside “loss.backward()” which you can use to accumulate the losses within the graph. The way ahead is to “step()” function once the model has seen enough batches of data. I have modified my training loop in which the model is stepped and then saving and loading the model again to resume the training. Although, It doesn’t look smooth, however, it worked for me.</p>
<p>So, setting the flag and modifying training process would solve the problem.</p>
<p>Hope this proves helpful!</p>
<p>Best regards</p>