pytorch/torch/csrc/autograd
Shivam Raikundalia 99c9a31386 [submodule] [Snapshot/Profiler] Memory Snapshot On Demand (#150559)
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
2025-04-07 13:04:38 +00:00
..
functions [ca] don't inline accumulate grad op (#149014) 2025-03-15 01:10:54 +00:00
utils [reland][ca] side-effect free inital trace: compiled_args (#148376) 2025-03-11 01:57:36 +00:00
anomaly_mode.cpp
anomaly_mode.h
autograd_meta.cpp
autograd_not_implemented_fallback.cpp
autograd_not_implemented_fallback.h
autograd.cpp
autograd.h
cpp_hook.cpp [ca] introduce RuntimeState to support c++ hooks via graph breaks (#149987) 2025-03-27 05:05:34 +00:00
cpp_hook.h [ca] introduce RuntimeState to support c++ hooks via graph breaks (#149987) 2025-03-27 05:05:34 +00:00
custom_function.cpp Improve error message when view of intermediate is returned from autograd.Function and marked dirty (#149543) 2025-03-25 14:44:11 +00:00
custom_function.h [reland][ca] side-effect free inital trace: compiled_args (#148376) 2025-03-11 01:57:36 +00:00
edge.h
engine.cpp [ca] support anomly mode nan checks with different semantics than eager (#149897) 2025-03-27 05:05:34 +00:00
engine.h [ca] side-effect free initial trace: GraphTask (#147796) 2025-02-26 16:37:27 +00:00
forward_grad.cpp
forward_grad.h
function_hook.h [reland][ca] side-effect free inital trace: compiled_args (#148376) 2025-03-11 01:57:36 +00:00
function.cpp
function.h [reland][ca] side-effect free inital trace: compiled_args (#148376) 2025-03-11 01:57:36 +00:00
FunctionsManual.cpp bf16 grouped gemm (#150374) 2025-04-06 04:53:24 +00:00
FunctionsManual.h bf16 grouped gemm (#150374) 2025-04-06 04:53:24 +00:00
grad_mode.h
graph_task.h
InferenceMode.h
init.cpp [custom_ops][perf] Move expensive pytree traversals of tensors to C++ (#148555) 2025-04-01 18:45:48 +00:00
input_buffer.cpp [MTIA] Ensure correct stream behavior for input_buffer add autograd on MTIA (#149433) 2025-03-19 20:19:18 +00:00
input_buffer.h
input_metadata.cpp
input_metadata.h
jit_decomp_interface.cpp
jit_decomp_interface.h
profiler_kineto.cpp [submodule] [Snapshot/Profiler] Memory Snapshot On Demand (#150559) 2025-04-07 13:04:38 +00:00
profiler_kineto.h [submodule] [Snapshot/Profiler] Memory Snapshot On Demand (#150559) 2025-04-07 13:04:38 +00:00
profiler_legacy.cpp
profiler_legacy.h
profiler_python.cpp [submodule] [Snapshot/Profiler] Memory Snapshot On Demand (#150559) 2025-04-07 13:04:38 +00:00
profiler_python.h
profiler.h
python_anomaly_mode.cpp
python_anomaly_mode.h
python_autograd.h
python_cpp_function.cpp Enable misc-use-internal-linkage check and apply fixes (#148948) 2025-03-12 14:22:56 +00:00
python_cpp_function.h
python_engine.cpp Enable misc-use-internal-linkage check and apply fixes (#148948) 2025-03-12 14:22:56 +00:00
python_engine.h
python_enum_tag.h
python_fft_functions.h
python_function.cpp Fix broken build within xplat/caffe2 (#149403) 2025-03-20 19:27:55 +00:00
python_function.h [reland][ca] side-effect free inital trace: compiled_args (#148376) 2025-03-11 01:57:36 +00:00
python_hook.cpp [reland][ca] side-effect free inital trace: compiled_args (#148376) 2025-03-11 01:57:36 +00:00
python_hook.h [reland][ca] side-effect free inital trace: compiled_args (#148376) 2025-03-11 01:57:36 +00:00
python_legacy_variable.cpp Enable misc-use-internal-linkage check and apply fixes (#148948) 2025-03-12 14:22:56 +00:00
python_legacy_variable.h
python_linalg_functions.h
python_nested_functions_manual.cpp Enable misc-use-internal-linkage check and apply fixes (#148948) 2025-03-12 14:22:56 +00:00
python_nested_functions.h
python_nn_functions.h
python_saved_variable_hooks.cpp [ca] trace saved variable unpacking (#147242) 2025-02-26 16:37:17 +00:00
python_saved_variable_hooks.h [ca] trace saved variable unpacking (#147242) 2025-02-26 16:37:17 +00:00
python_sparse_functions.h
python_special_functions.h
python_torch_functions_manual.cpp Enable misc-use-internal-linkage check and apply fixes (#148948) 2025-03-12 14:22:56 +00:00
python_torch_functions.h Enable misc-use-internal-linkage check and apply fixes (#148948) 2025-03-12 14:22:56 +00:00
python_variable_indexing.cpp
python_variable_indexing.h
python_variable.cpp Enable misc-use-internal-linkage check and apply fixes (#148948) 2025-03-12 14:22:56 +00:00
python_variable.h
README.md
record_function_ops.cpp Enable misc-use-internal-linkage check and apply fixes (#148948) 2025-03-12 14:22:56 +00:00
record_function_ops.h
saved_variable_hooks.h [ca] trace saved variable unpacking (#147242) 2025-02-26 16:37:17 +00:00
saved_variable.cpp [ca] trace saved variable unpacking (#147242) 2025-02-26 16:37:17 +00:00
saved_variable.h [ca] trace saved variable unpacking (#147242) 2025-02-26 16:37:17 +00:00
symbolic.h
TraceTypeManual.cpp
variable_info.cpp
variable_info.h
variable.cpp
variable.h
VariableTypeManual.cpp
VariableTypeUtils.h

Autograd

Autograd is a hotspot for PyTorch performance, so most of the heavy lifting is implemented in C++. This implies that we have to do some shuffling between Python and C++; and in general, we want data to be in a form that is convenient to manipulate from C++.

Our general model is that for any key data type that autograd manipulates, there are two implementations: a C++ type and a Python object type. For example, consider variables in autograd: we have both Variable in variable.h (the C++ type) and THPVariable in python_variable.h (the Python type.) (By the way, THP stands for TorcH Python, not to be confused with THPP, TorcH C++). Variable contains the payload of a variable, while THPVariable just contains a shared_ptr reference to Variable, as well as references to other Python objects which the Python runtime needs to know about. A lot of data accessor implementations in python_variable.cpp simply reach through to the underlying Variable and return the appropriate value.

The most complicated application of this principle is Function, which also supports users implementing custom behavior in Python. We have the following classes:

  • Node in function.h, the C++ type.
  • THPFunction in python_function.h, the Python object type. In python_function.cpp, you can see the boilerplate that tells the Python interpreter about this object.
  • PyNode in python_function.h, a subclass of Node which forwards apply to a Python THPFunction. (NOT a Python object, despite its name!)

Outside of PyNode, the C++ objects largely avoid referencing Python objects (there are a few exceptions, like pyobj in Variable, and PyNode, whose whole point is to let C++ call into Python). And pyobj in Node to ensure uniqueness of the associated python wrapper (if it exists).