How to integrate the profile tool into the transformers framework

[Hukongtao]

Description & Motivation

First of all, thank you for your contribution.
I am training a model using the transformers framework, but I want to use your profile tool.
Is there any quick way to do this?

Pitch

No response

Alternatives

No response

Additional context

No response

[SkafteNicki]

Hi, thanks for raising this issue. It is fairly simple to integrate lightnings profilers with transformers framework or for that sake any other framework that you may want. There are two steps: initialize profiler

from lightning.pytorch.profilers import SimpleProfiler
profiler = SimpleProfiler(
        dirpath="./profiler_results",      # Where to save results
        filename="transformers_profile",   # File name
        extended=True                       # Show detailed stats
)

use profiler.profile as context manager around what you want to profile

with profiler.profile():
    # code to profile

a full example using transformers:

from lightning.pytorch.profilers import SimpleProfiler
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
from datasets import Dataset
import torch

def create_sample_dataset():
    """Create a simple dataset for demonstration."""
    texts = ["I love this!", "This is terrible", "Not bad", "Amazing!"] * 25
    labels = [1, 0, 0, 1] * 25  # 1=positive, 0=negative
    return Dataset.from_dict({"text": texts, "labels": labels})

def main():
    profiler = SimpleProfiler(
        dirpath="./profiler_results",      # Where to save results
        filename="transformers_profile",   # File name
        extended=True                       # Show detailed stats
    )
    
    with profiler.profile("model_tokenizer_loading"):
        model_name = "distilbert-base-uncased"
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
    
    with profiler.profile("dataset_preparation"):
        dataset = create_sample_dataset()
        
        def tokenize_function(examples):
            return tokenizer(examples["text"], truncation=True, padding=True, max_length=128)
        
        dataset = dataset.map(tokenize_function, batched=True)
        train_dataset = dataset.train_test_split(test_size=0.2)["train"]
    
    with profiler.profile("training_setup"):
        training_args = TrainingArguments(
            output_dir="./results",
            num_train_epochs=1,
            per_device_train_batch_size=8,
            logging_steps=10,
            save_strategy="no",
        )
        
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=train_dataset,
            tokenizer=tokenizer,
        )
    
    with profiler.profile("training_execution"):
        trainer.train()
    
    with profiler.profile("evaluation"):
        eval_dataset = dataset.train_test_split(test_size=0.2)["test"]
        results = trainer.evaluate(eval_dataset)
        print(f"Evaluation results: {results}")
    
    print(profiler.summary())
    
    profiler.describe()


if __name__ == "__main__":
    main()

you can replace the simple profiler with any of the other profilers that lightning comes with:
https://lightning.ai/docs/pytorch/stable/api_references.html#profiler
Closing issue, but feel free to ping me and reopen if necessary.