UPMEM LLM framework for profiling / simulation
A Python library that profiles the execution time of individual PyTorch layers and functions on a CPU, then simulates how those same layers would run and how much energy they would use on a target hardware accelerator profile.
The library is imported alongside Typer to build a small profiling CLI: a callback initializes the library's profiling options, and profiler_init, profiler_start and profiler_end calls bracket a model's forward pass to record timing. Each layer or function the profiler tracks is identified by four things: its type, such as a Linear module, or the function name such as softmax; the variable name it was assigned to when it was created; its input dimensions; and, for a layer, a unique ID assigned at initialization.
By default the profiler reports a summary of execution time at the end of a run, and when simulating a target hardware accelerator it also reports estimated energy and power consumption, broken down between the summarization or encoding phase and the generation phase; it does not currently track the CPU's real power draw during plain profiling. Usage examples are included for both a custom PyTorch model and a model loaded from Hugging Face.
It is written in Python.