99c9a31386
Summary:
Profiler side of memory snapshot.
1. Add API to actually do snapshot when client interface is called
2. Add ifdefs to builds so that kineto hooks snapshot correctly.
Design Philosophy: There is one interesting part of this implementation and it is during export. For export we are callign the python impl of the export rather than CPP even though we are already in CPP. This is because it is better to simply have one path of export rather than 2. Personally, I want there to be parity between auto-trace and on-demand so it if we can limit the side paths then we will have an easier time maintaining this relationship
Test Plan: {F1976563426}
Reviewed By: sanrise
Differential Revision: D70733247
Pull Request resolved: https://github.com/pytorch/pytorch/pull/150559
Approved by: https://github.com/sanrise
|
||
|---|---|---|
| .. | ||
| orchestration | ||
| python | ||
| standalone | ||
| stubs | ||
| unwind | ||
| api.h | ||
| collection.cpp | ||
| collection.h | ||
| combined_traceback.cpp | ||
| combined_traceback.h | ||
| containers.h | ||
| data_flow.cpp | ||
| data_flow.h | ||
| events.h | ||
| kineto_client_interface.cpp | ||
| kineto_client_interface.h | ||
| kineto_shim.cpp | ||
| kineto_shim.h | ||
| perf-inl.h | ||
| perf.cpp | ||
| perf.h | ||
| README.md | ||
| util.cpp | ||
| util.h | ||
Profiler Overview
This README describes the details of how the profiler is implemented.
The profiler instruments PyTorch to collect information about the model's execution. Its main features are:
- Instrumenting op calls on the CPU side
- Interfacing with Kineto to collect information from the GPU (or other accelerators)
- Collecting python stack traces
- Exporting this information, e.g. in a chrome trace, or to be processed by downstream tools like HTA
Table of Contents
- Codebase Structure
RecordFunction- Autograd Integration
- Collection and Post-Processing
- Kineto Integration
- Python Tracing
Codebase Structure
TODO
RecordFunction
/aten/src/ATen/record_function.h
RecordFunction is used by the profiler to instrument CPU-side events.
RecordFunction is a general method of instrumenting function calls in PyTorch. It can be used for other general applications, e.g. see Features for Large-Scale Deployments. In PyTorch, it is already included at some important locations; notably, in the dispatcher, surrounding every op.
Users (or PyTorch itself) can register callbacks that will be executed whenever a RecordFunction guard is encountered. The profiler uses this mechanism to record the start and end times for each op call, as well as user-provided RecordFunction annotations. The RecordFunction machinery is designed to have relatively low overhead, especially when there are no callbacks registered. Nevertheless, there can still be some overhead.
There is also a python binding for RecordFunction in python (with torch.profiler.record_function); this is often used by users to annotate events corresponding to module-level events.
Autograd Integration
The autograd engine is responsible for automatically computing gradients.
The profiler records two pieces of information from the autograd engine:
- Sequence number: this is a unique-per-thread index assigned to each op call(*) in the forward pass. When a backward op is triggered, it is also assigned a sequence number matching the sequence number of the forward op that caused that backward op to be executed. Using this information, the profiler is able to match forward and backward ops; in chrome traces, this feature can be enabled with the "fwd_bwd" flow events
- Forward thread id: Autograd can be used in multi-threaded environments. The forward thread ID indicates the ID of the thread on which the forward op was executed on. This information is needed because the sequence number, mentioned above, is only unique within a thread; the forward thread ID is used for differentiating different ops with the same sequence number.
(*) Note that only op invocations whose inputs require gradients are assigned a sequence number
Collection and Post-Processing
TODO
Kineto Integration
TODO
Python Tracing
TODO