Is there a pytorch profiler integration with huggingface trainer?

sajaldash · August 31, 2022, 6:49pm

I am trying to profile various resource utilization during training of transformer models using HuggingFace Trainer. Since the HF Trainer abstracts away the training steps, I could not find a way to use pytorch trainer as shown in here. I can extend the HF Trainer class and overwrite the train() function to integrate the profiler.step() instruction, but the train() function is a lengthy and complex one. Is there an easy way to using PyTorch Profiler with HuggingFace Trainer class?

Data-drone · March 17, 2023, 2:25am

It looks like this an example for using profiler with huggingface:

github.com

yqhu/profiler-workshop/blob/c8d4a7c30a61cc7b909d89f88f5fd36b70c55769/hf_training_trainer_prof.py

from datasets import load_dataset, load_metric
from transformers import AutoTokenizer, DataCollatorWithPadding, AutoModelForSequenceClassification
from transformers import Trainer, TrainingArguments, TrainerCallback
import torch
import numpy as np
import time


raw_datasets = load_dataset("glue", "mrpc")
checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

def tokenize_function(example):
    return tokenizer(example["sentence1"], example["sentence2"], truncation=True)

tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)

This file has been truncated. show original