After 50-100 Passes, the Model Runs Out of Memory

vsiddons: <p>I am currently attempting to fine-tune Salesforce/Blip-2-opt-2.7b using the HuggingFace transformers library. I am training a LoRA adapter using peft.</p>
<p>My problem: After I begin training the model, depending on the batch size, my model will run out of memory after ~50 (batch size of 32) or ~100 (batch size of 8) passes. It seems that something is slowly eating up memory, but I cannot tell.</p>
<p>Here is the relevant code:</p>
<pre><code class="lang-auto">model.train()
        print("Starting...\n")
        for epoch in range(0, epochs):
            for sample in data_loader:

                #responces = [requests.get(x) for x in sample["image"]]
                images = [Image.open(x) for x in sample["image"]]

                #images = sample["image"]

                #https://github.com/huggingface/notebooks/blob/main/peft/Fine_tune_BLIP2_on_an_image_captioning_dataset_PEFT.ipynb
                encoded_images = torch.Tensor(len(images), 3, 224, 224)
                encoded_images = encoded_images.to(device)
                for i in range(0, len(images)):
                    #print(processor(images=images[i], padding="max_length", return_tensors="pt"))
                    image = processor(images=images[i], padding="max_length", return_tensors="pt")["pixel_values"]
                    image = torch.squeeze(image)
                    image = image.to(device)
                    encoded_images[i] = image
                # remove batch dimension
                #encoded_images = torch.squeeze(encoded_images)

                for i in range(0, len(sample["t"])):
                    sample["t"][i] = str(sample["t"][i])
                for i in range(0, len(sample["s"])):
                    sample["s"][i] = str(sample["s"][i])

                t_questions = []
                s_questions = []
                
                counter = 0
                for caption in sample["caption"]:
                    t_questions.append(str(caption) + " Question: What is used as for " + sample["s"][counter] + "? Answer: " + 
                                            sample["t"][counter])
                    s_questions.append(str(caption) + " Question: What is " + sample["t"][counter] + "? Answer: " + 
                                            sample["s"][counter])
                    counter += 1             

                processed_tq = torch.LongTensor(processor.tokenizer(t_questions, padding= True, max_length= 64)["input_ids"])
                processed_tq = processed_tq.to(device)

                #with autocast():
                    #t_output = model(input_ids= processed_tq, pixel_values= encoded_images, labels= processed_tq)
                    #loss = t_output.loss.mean()

                #scaler.scale(loss).backward()

                t_output = model(input_ids= processed_tq, pixel_values= encoded_images, labels= processed_tq)

                loss = t_output.loss.mean()

                print("T Loss: " + str(loss.item()) + "\n")

                #torch.cuda.empty_cache()
                    
                loss.backward()
                #scaler.step(optimizer)
                #scaler.update()
                #optimizer.zero_grad()

                optimizer.step()
                optimizer.zero_grad()

                ##ty_cache()
                del processed_tq, t_output
                gc.collect()

                processed_sq = torch.LongTensor(processor.tokenizer(s_questions, padding= True, max_length= 64)["input_ids"])
                processed_sq = processed_sq.to(device)

                #with autocast():
                    #s_output = model(input_ids= processed_sq, pixel_values= encoded_images, labels= processed_sq)
                    #loss = s_output.loss.mean()
                
                #scaler.scale(loss).backward()

                #scaler.step(optimizer)
                #scaler.update()
                #optimizer.zero_grad()

                s_output =  model(input_ids= processed_sq, pixel_values= encoded_images, labels= processed_sq)

                loss = s_output.loss.mean()

                print("S Loss: " + str(loss.item()) + "\n")

                #torch.cuda.empty_cache()

                loss.backward()

                optimizer.step()
                optimizer.zero_grad()

                #torch.cuda.empty_cache()
                del processed_sq, s_output
                gc.collect()

                for image in images:
                    image.close()
                
                #torch.cuda.empty_cache()
                del images, encoded_images
                gc.collect()

            epochs_list.append(epoch)
            losses.append(loss.mean().item())
        print("Finished...\n")
</code></pre>
<p>Things I have tried so far that haven’t worked:</p>
<ul>
<li>I tried reducing the batch size and the lora rank.</li>
<li>To remedy this, I tried using normal data parallel, but it calculates the t loss and then gives up.</li>
<li>I then tried to use DistributedDataParallel but when I entered an empty port, the model just gets stuck on loading itself.</li>
<li>Note: All of this worked fine on a similar dataset with a batch size of 32. I never had the same problem.</li>
</ul>