Как я могу группировать события PyTorch Profiler по иерархии слоев при профилировании Hugging Face Transformer? ⇐ Python

[Anonymous]

Я использую PyTorch Profiler для проверки производительности вывода на преобразователе Hugging Face (например, модели Qwen). У меня есть код, который успешно собирает информацию профилирования на уровне оператора (например, aten::mm, aten::copy_, aten::pow и т. д.), но трудно понять, какая именно. слое или подмодуле, которому принадлежит каждый вызов оператора.
Вопрос:

• Есть ли встроенный или сторонний способ автоматической группировки событий профилировщика по подмодулю или уровень, который их вызвал?
• Если нет, то какой подход рекомендуется (например, ручное инструментирование с помощью Record_function, обезьянье исправление или что-то еще) для маркировки или группировки этих событий по иерархию слоев в результатах профилирования?

Я читал о ручном вызове функции Record_function("layer_name") и об исправлении обезьян вперед( ), но я не уверен, что существует более стандартное или библиотечное решение специально для трансформеров Hugging Face Transformers.
Будем очень признательны за любые советы и рекомендации!
Вот упрощенная версия моей текущей настройки:

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

import torch
from torch.profiler import profile, ProfilerActivity
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_NAME = "Qwen/Qwen2.5-0.5B"
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16, device_map="auto")
model.eval()

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

prompt = "Please introduce the US EV company Tesla."
inputs = tokenizer(prompt, return_tensors="pt")
input_device = next(model.parameters()).device
for k, v in inputs.items():
inputs[k] = v.to(input_device)

def generate_text_and_profile(model, inputs, max_new_tokens=64):
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
record_shapes=True,
profile_memory=True,
with_stack=False,
with_flops=True,
) as prof:
generated_ids = model.generate(**inputs, max_new_tokens=max_new_tokens)

print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=-1))
prof.export_chrome_trace("trace.json")
return generated_ids

generated_ids = generate_text_and_profile(model, inputs)
decoded = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
print(decoded)

Таблица профилировщика показывает стоимость каждого оператора (например, aten::mm, aten::copy_), но не указывает, какая часть модели (например, «LayerNorm в блоке 1», «Multi-Head Attention в блоке 3» или «FeedForward в блоке 7»), к которой принадлежат эти операторы.

Подробнее здесь: https://stackoverflow.com/questions/793 ... ing-a-hugg