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#environment-variables)

Comments

[volosogriz]

Сменил модель на facebook/xglm-564M. Она начинает тренировку, но появляется такая же ошибка на 32%. Если изменить batch size, то ошибка будет сразу же.

Answer 1

[rPman]

Твоя видеокарта всего 4гб, откуда ты взял 7гб всего?
Скорее всего твой датасет не помещается в память видеокарты... уменьши размер батча (вообще то его увеличивают чтобы значительно ускорить обучение, но тебе 'не до жиру') до 1.
Перейди на fp16, это уменьшит размер весов модели (правда если она уже уменьшена, дистиллирована и прочее, это может помешать обучению)

training_args = TrainingArguments(
    fp16=True,

Так же можно поделить свою обучающую выборку на части и обучать сетку по ее частям (это замедляет процесс и ухудшает но все еще лучше чем на процессоре все делать)

training_args = TrainingArguments(
    gradient_accumulation_steps=4,