Torch.OutOfMemoryError PyTorch занимает много памяти. Как оптимизировать?

Question

[volosogriz]

Доброго времени суток! Столкнулся с проблемой при попытке запустить тренировку mbart many-to-many на rtx 3050 ti с 4гб выделенной и 7гб общей памяти. PyTorch почему-то съедает всю видеопамять.
Возможно я чего-то недопонимаю. Буду рад любым советам!

import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
import torch
from torch.utils.data import Dataset
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
import pandas as pd
from torch.utils.data import DataLoader
from transformers import Trainer, TrainingArguments

model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

tokenizer.pad_token = tokenizer.eos_token

class TextDataset(Dataset):
    def __init__(self, encodings):
        self.input_ids = encodings['input_ids']
        self.attention_mask = encodings['attention_mask']
        self.labels = encodings['input_ids'].clone()

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, idx):
        return {
            'input_ids': self.input_ids[idx],
            'attention_mask': self.attention_mask[idx],
            'labels': self.labels[idx], 
        }

@dp.message_handler(commands=['train'])
async def get_users_handler(message: types.Message):

    # dataset
    df = pd.read_csv("dataset.csv")
    if 'Вопрос' not in df.columns or 'Ответ' not in df.columns:
        return
    texts = df['Вопрос'] + " " + df['Ответ']

    encodings = tokenizer(texts.tolist(), truncation=True, padding=True, return_tensors='pt')
    dataset = TextDataset(encodings)

    dataloader = DataLoader(dataset, batch_size=2)  # batch 

    for batch in dataloader:
        # Перемещаем данные на устройство
        inputs = {
            'input_ids': batch['input_ids'].to(device),
            'attention_mask': batch['attention_mask'].to(device),
            'labels': batch['labels'].to(device)
        }
    
        outputs = model(**inputs)
        loss = outputs.loss
        print(loss)

    # args
    training_args = TrainingArguments(
        output_dir='./results',
        num_train_epochs=3,
        per_device_train_batch_size=2,
        save_steps=10_000,
        save_total_limit=2,
        logging_dir='./logs',  
        logging_steps=500,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
    )

    trainer.train()

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 20.00 MiB. GPU 0 has a total capacity of 4.00 GiB of which 0 bytes is free. Of the allocated memory 10.49 GiB is allocated by PyTorch, and 123.19 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#en...)