Загрузка графического процессора во время тренировки почти всегда 0. Hugging Face Transformer ⇐ Python

[Anonymous]

Я настраиваю модель Donut Cord-v2 с использованием данных своего счета, размер которых после предварительной обработки и сохранения на диске в виде набора данных составляет около 360 ГБ. Я почти точно следую этому ноутбуку, за исключением того, что у меня 6 эпох обучения вместо 3.
Я тренируюсь на одном графическом процессоре Nvidia H100 SXM / Intel Xeon® Gold 6448Y / 128 ГБ ОЗУ.< /p>
Всякий раз, когда я начинаю обучение и проверяю загрузку ЦП и ГП с помощью htop и nvidia-smi, я вижу, что ЦП находится на уровне 10–12 %. использование, используется Python, память графического процессора постоянно заполнена почти на 90%, но загрузка графического процессора почти всегда равна 0. Если я продолжу обновлять вывод nvidia-smi, раз в 10-12 секунд загрузка подскочит до 100 %, а затем немедленно вернуться к 0. Я не могу не чувствовать, что между моим процессором и графическим процессором существует узкое место, когда процессор пытается постоянно обрабатывать данные и отправлять их в графический процессор, графический процессор обрабатывает их очень быстро и просто простаивает, ожидая следующего пакета от процессора. Я загружаю уже предварительно обработанный набор данных с диска следующим образом:

Код: Выделить всё

from datasets import load_from_disk
processed_dataset = load_from_disk(r"/dataset/dataset_final")

Моя конфигурация процессора следующая:

Код: Выделить всё

from transformers import DonutProcessor

new_special_tokens = [] # new tokens which will be added to the tokenizer
task_start_token = ""  # start of task token
eos_token = "" # eos token of tokenizer

processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")

# add new special tokens to tokenizer
processor.tokenizer.add_special_tokens({"additional_special_tokens": new_special_tokens + [task_start_token] + [eos_token]})

# we update some settings which differ from pretraining; namely the size of the images + no rotation required
processor.feature_extractor.size = [1200,1553] # should be (width, height)
processor.feature_extractor.do_align_long_axis = False

Конфигурация моей модели:

Код: Выделить всё

import torch
from transformers import VisionEncoderDecoderModel, VisionEncoderDecoderConfig

#print(torch.cuda.is_available())

# Load model from huggingface.co
model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")

# Resize embedding layer to match vocabulary size
new_emb = model.decoder.resize_token_embeddings(len(processor.tokenizer))
print(f"New embedding size: {new_emb}")
# Adjust our image size and output sequence lengths
model.config.encoder.image_size = processor.feature_extractor.size[::-1] # (height, width)
model.config.decoder.max_length = len(max(processed_dataset["train"]["labels"], key=len))

# Add task token for decoder to start
model.config.pad_token_id = processor.tokenizer.pad_token_id
model.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids([''])[0]

И мой код обучения:

Код: Выделить всё

import gc
gc.collect()

torch.cuda.empty_cache()

from transformers import Seq2SeqTrainingArguments, Seq2SeqTrainer

import logging
logging.basicConfig(level=logging.INFO)

# Arguments for training
training_args = Seq2SeqTrainingArguments(
output_dir=r"/trained",  # Specify a local directory to save the model
num_train_epochs=6,
learning_rate=2e-5,
per_device_train_batch_size=8,
weight_decay=0.01,
fp16=True,
logging_steps=50,
save_total_limit=2,
evaluation_strategy="no",
save_strategy="epoch",
predict_with_generate=True,
report_to="none",
# Disable push to hub
push_to_hub=False

)

# Create Trainer
trainer = Seq2SeqTrainer(
model=model,
args=training_args,
train_dataset=processed_dataset["train"],
)

# Start training
trainer.train()

Оценочное время завершения обучения с использованием 6 эпох и набора данных объемом 360 ГБ составляет 54 часа. Когда я запускаю тот же самый код на своем компьютере с Intel i9 11900KF / RTX 3050, я постоянно вижу загрузку графического процессора на уровне 100%. Есть ли узкое место в моем коде? Почему процессор продолжает так много обрабатывать уже предварительно обработанный набор данных? Cuda 12.6
Изменить:
Имеет ли смысл изменить параметр dataloader_num_workers Seq2SeqTrainer на значение >0, поскольку моя оперативная память и Количество ядер процессора позволяет? (и поскольку загрузка ЦП составляет максимум 10–12 %)

Подробнее здесь: https://stackoverflow.com/questions/793 ... ransformer