PyTorch Profiler 實現原理深度解析從 RecordFunction 到 Kineto 的 CPU/GPU 全鏈路采集架構【免費下載鏈接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration項目地址: https://gitcode.com/GitHub_Trending/py/pytorchtorch.profiler是 PyTorch 官方提供的性能剖析工具它能夠記錄模型執行過程中的算子耗時、內存分配、Python 調用棧與 GPU 事件并導出為 Chrome Trace / Perfetto 可讀的 JSON 文件。本文以倉庫中 torch/csrc/profiler/README.md 為核心骨架結合 torch/csrc/profiler/collection.h、torch/csrc/profiler/collection.cpp、torch/csrc/autograd/profiler_kineto.cpp 等源碼深入剖析 Profiler 的底層實現事件從 CPU 算子回調、內存分配、自動求導、Python 棧到 GPU 活動是如何被采集、關聯與導出的。讀完本文你將理解 Profiler 的完整數據鏈路、各采集階段的內部機制以及如何基于這些原理正確配置和使用剖析參數。代碼庫結構Profiler 前端、后端與 C 采集層的職責劃分Profiler 的實現橫跨 Python 與 C 兩大層次README 給出了核心文件布局相對倉庫根路徑關鍵組成如下torch/ │ ├── profiler/ # 主 Python 包包含核心前端邏輯 │ ├── __init__.py # profiler 包的初始化文件 │ ├── profiler.py # 主 Profiler 前端類profile、schedule、supported_activities 等 │ └── _utils.py # FunctionEvent 工具函數 │ ├── autograd/ # Autograd 包 │ ├── __init__.py # autograd 包初始化文件 │ ├── profiler.py # 主 Profiler 后端類 │ └── profiler_utils.py # FunctionEvent 工具函數 │ ├── csrc/ # C 與 C 源碼 │ └── profiler/ # Profiler C 源碼 │ ├── collection.cpp # 主要采集邏輯 │ ├── collection.h # 采集定義 │ ├── kineto_client_interface.cpp # 從 kineto 調用 Profiler 的接口僅 on-demand 模式 │ ├── kineto_client_interface.h # Client 接口定義 │ ├── kineto_shim.cpp # 從 profiler 調用 kineto 的 shim │ ├── kineto_shim.h # shim 定義 │ ├── util.cpp # 處理 profiler 事件中參數的 utils │ ├── util.h # util 定義 │ └── README.md # 本文所講解的文檔 │ └── autograd/ # Autograd C 源碼 │ ├── profiler_python.cpp # 主要的 Python 調用棧采集邏輯 │ ├── profiler_python.h # Python 調用棧采集定義 │ ├── profiler_kineto.cpp # 啟動采集/kineto 的 Profiler 后端邏輯 │ └── profiler_kineto.h # 啟動采集/kineto 的 Profiler 后端定義 │ └── ATen/ # ATen C 源碼 │ ├── record_function.cpp # RecordFunction 采集邏輯 │ └── record_function.h # RecordFunction 定義從源碼結構可以清晰地看到職責邊界Python 前端torch/profiler/profiler.py負責向用戶暴露profile()、schedule()、ProfilerActivity等 API并把用戶的配置如record_shapes、with_stack翻譯成底層ProfilerConfig。Python 后端torch/autograd/profiler.pytorch.autograd.profiler模塊中的profile類歷史上是舊式 Profiler 的入口現在與 Kineto 后端并存。C 采集層torch/csrc/profiler/實現線程級事件緩沖ThreadLocalSubqueue、事件樹構建Result、Kineto shim 與時鐘轉換。Autograd C 層torch/csrc/autograd/profiler_kineto.cpp實現回調注冊與前后端橋接profiler_python.cpp實現基于 CPython C API 的調用棧追蹤。RecordFunctionCPU 側事件插樁的基礎設施RecordFunction是 Profiler 插樁 CPU 側事件的核心機制其定義位于 aten/src/ATen/record_function.h。它是一個通用的函數調用插樁手段并非 Profiler 專有也可以用于其他通用場景例如 PyTorch 的 大規模部署特性。在 PyTorch 內部它已被安插在若干關鍵位置最典型的是dispatcher中環繞每一個算子調用見 aten/src/ATen/core/dispatch/Dispatcher.h。RecordFunction的核心設計是回調注冊制用戶或 PyTorch 自身可以注冊回調每當執行路徑遇到一個RecordFunction守衛時這些回調就會被觸發。Profiler 正是利用這一機制來記錄每個算子調用的開始與結束時間以及用戶自定義的RecordFunction注釋。從工程角度RecordFunction機制被刻意設計為低開銷尤其是在沒有注冊任何回調時。這是因為算子調度路徑是訓練/推理的熱路徑任何額外開銷都會被放大。盡管如此一旦啟用回調依然會引入一定的性能損耗——這解釋了為什么 Profiler 只在需要時開啟采集。RecordFunction還提供了 Python 綁定with torch.profiler.record_function(...)。這是用戶在代碼中標記模塊級事件的常用手段例如import torch with torch.profiler.profile(activities[torch.profiler.ProfilerActivity.CPU]) as prof: with torch.profiler.record_function(my_custom_module): x torch.randn(1024, 1024) y x.mm(x)自定義的USER_SCOPE事件與算子事件一樣會被采集并出現在最終的 trace 中。Autograd 集成用序列號與前向線程 ID 關聯前反向自動求導引擎負責自動計算梯度。Profiler 從 autograd 引擎記錄兩類關鍵信息用于在 trace 中把反向算子與觸發它的前向算子關聯起來序列號Sequence Number定義位于 aten/src/ATen/SequenceNumber.h。這是一個每線程唯一的索引分配給前向傳播中的每次算子調用。當反向算子被觸發時它會被賦予與其來源前向算子相同的序列號。利用這一點Profiler 能夠完成前向/反向算子的匹配在 Chrome Trace 中該功能以fwd_bwd flow events的形式呈現。需要特別注意的是README 中的腳注只有輸入張量需要梯度的算子調用才會被分配序列號。這意味著純推理無梯度需求場景下序列號關聯的價值有限而在訓練場景中該機制是剖析反向傳播熱點的重要依據。前向線程 IDForward Thread IDAutograd 可以在多線程環境中使用。前向線程 ID 表示前向算子執行所在線程的 ID相關字段可見于 aten/src/ATen/record_function.h 的RecordFunction定義中。之所以需要它是因為上述序列號只在單個線程內唯一不同線程上可能產生相同的序列號必須用前向線程 ID 加以區分。在 torch/csrc/profiler/collection.h 的TorchOpBasicFields中可以看到這兩個字段的直接體現struct TorchOpBasicFields { int64_t sequence_number_{0}; uint64_t forward_tid_{0}; ... };ExtraFieldsEventType::TorchOp繼承TorchOpBasicFields在采集時隨事件一起寫入供后處理階段構建前反向關聯。Torch 算子采集流程回調注冊、線程子隊列與熱路徑優化本節描述auto-trace進程內、同步模式下 torch 算子的通用采集流程。關于 on-demand進程外、異步追蹤的細節可參考 Libkineto 的 READMEKineto 作為第三方子模塊被引入。回調的注冊profiler_kineto.cpp當一次 trace 開始時autograd/profiler 后端會調用 torch/csrc/autograd/profiler_kineto.cpp 來準備、啟動或停止采集。在 trace 啟動時定義于該文件中的onFunctionEnter與onFunctionExit回調會被注冊到RecordFunction機制中。源碼中回調分為兩組torch/csrc/autograd/profiler_kineto.cpponFunctionEnterGlobal/onFunctionExitGlobal全局回調用于KINETO_ONDEMAND或配置了profile_all_threads的KINETO會話onFunctionEnterTLS/onFunctionExitTLS線程本地回調僅對 trace 啟動時存在的線程生效。回調的注冊范圍由ExperimentalConfig決定分為兩種模式Global全局回調注冊到執行期間的所有線程。Local本地回調僅注冊到trace 開始時刻已存在的線程上。auto recordFunctionCallback at::RecordFunctionCallback(onFunctionEnterGlobal, onFunctionExitGlobal) .needsInputs(state_ptr-config().report_input_shapes) .scopes(scopes);needsInputs(config.report_input_shapes)表明只有當用戶開啟record_shapes時回調才需要算子輸入信息否則采集開銷更小。線程子隊列與begin_op在onFunctionEnter內部Profiler 會為每個線程創建一個ThreadLocalSubqueue實例見 torch/csrc/profiler/collection.h確保每個 CPU 算子都與它實際執行的線程關聯。當一個 torch 算子進入時Profiler 調用定義在collection.cpp的begin_op來記錄必要信息。ThreadLocalSubqueue的設計值得關注內部使用AppendOnlyList只追加、塊大小為 512作為事件存儲避免每次算子都分配獨立 vector每種事件類型對應獨立的存儲區torch_ops_算子事件、allocations_內存分配、backend_events_后端事件、vulkan_events_、ooms_OOM、py_calls_Python 調用算子事件的correlation id由EventBlock按塊起始 ID 塊內偏移計算無需額外分配torch/csrc/profiler/collection.cpp。README 特別強調begin_op被刻意設計得非常輕量因為它處于 profiling 的 hot path熱路徑上過高的額外開銷會扭曲 profile 結果、降低其參考價值。因此回調期間只收集最必要的信息大部分邏輯都推遲到后處理階段完成。這也是InputOutputEncoder存在的原因——它在采集期把算子輸入的形狀、dtype 編碼進連續的AppendOnlyList避免為每個算子創建 vector而在后處理時才解碼還原見 torch/csrc/profiler/collection.cpp。全局會話的并發安全對于全局回調profiler_kineto.cpp中還有一個GlobalCallbackSession協調機制torch/csrc/profiler/kineto_profiler.cpp由于全局回調會在任意線程上觸發而disableProfiler()可能在另一個線程上銷毀 profiler 狀態因此每個回調通過enter()/exit()只在極短的臨界區getGlobal()RecordQueue訪問內持有 in-flight 計數teardown 時drain()關閉會話并等待所有 in-flight 回調退出避免 use-after-free。退出回調還會校驗session_generation_丟棄跨會話懸空的退出事件。內存分配事件采集繞過 RecordFunction 的瞬時事件與有起止時刻的算子事件不同內存分配事件被表示為cpu_instant_event零持續時間。因此分配事件不經過RecordFunction而是由reportMemoryUsage直接調用emplace_allocation_event把事件入隊到對應的ThreadLocalSubqueue。在 torch/csrc/autograd/profiler_kineto.cpp 中可以看到void reportMemoryUsage( void* ptr, int64_t alloc_size, size_t total_allocated, size_t total_reserved, c10::Device device) override { if (config_.profile_memory !config_.disabled()) { recordQueue.getSubqueue()-emplace_allocation_event( c10::getApproximateTime(), ptr, alloc_size, total_allocated, total_reserved, device.type(), device.index()); } }同樣地OOM內存不足事件通過reportOutOfMemory→emplace_ooms_event進入隊列。RawAllocation結構torch/csrc/profiler/collection.h記錄了分配指針、大小、累計分配量/預留量以及設備信息并被static_assert(std::is_trivial_vRawAllocation)強制保持平凡類型以提升性能。在 Python 前端對應的開關是profile()的profile_memory參數torch/profiler/profiler.pytorch.profiler.profile( activities[torch.profiler.ProfilerActivity.CPU], profile_memoryTrue, # 跟蹤張量內存分配/釋放 )Kineto 集成多架構事件采集、預熱與導出Kineto 是 Profiler 獲取 GPU 及其他加速器事件的關鍵抽象層。它作為第三方子模塊third_party/kineto/目錄被引入通過與 CUPTI 等庫交互來接收 GPU 與加速器事件再轉發給前端 Profiler。預熱warmup的原因Kineto 需要時間prepare也稱 warmup這些第三方模塊以避免初始化例程扭曲 profile 結果。理論上可以在作業啟動時就完成預熱但讓 CUPTI 這類重量級庫持續運行會帶來顯著的不必要開銷因此預熱被安排在正式 trace 之前的短暫窗口內完成。雙向橋接與后處理如前所述profiler_kineto.cpp在后臺調用合適的 profiler 階段同時它還會調用kineto_shim.cpp后者觸發 Kineto 中的對應例程。當一次 trace 完成后Kineto 收集的所有事件會被轉發給 Profiler原因有兩個合并所有數據并完成 Profiler 事件與 Kineto 事件之間的后處理例如把 Kineto 事件嵌入Result事件樹見 torch/csrc/profiler/collection.h 的ExtraFieldsEventType::Kineto其中Flow結構用于將 Kineto 的 flow 事件映射進 profiler 樹把這些事件作為FunctionEvents轉發給 Python 前端供prof.key_averages()、prof.table()等 API 使用。文件導出集成的最后一步是文件導出。所有事件收集并后處理完成后它們可以導出為 JSON 文件供Perfetto或Chrome Tracer可視化。這一步通過調用 Kineto 的ActivityTraceInterface::save完成把所有事件信息寫入磁盤。在 Python 側tensorboard_trace_handler會自動生成 trace 目錄from torch.profiler import profile, ProfilerActivity, tensorboard_trace_handler with profile( activities[ProfilerActivity.CPU, ProfilerActivity.CUDA], on_trace_readytensorboard_trace_handler(./log/resnet18), ) as prof: model(x) # 輸出的 JSON 文件位于 ./log/resnet18/ 目錄下Python 棧追蹤基于 CPython 剖析 API 的實現當在 profiler 中設置with_stackTrue時Python 棧追蹤器會使用PythonTracerBase中定義的make函數生成具體實現位于profiler_python.cpptorch/csrc/autograd/profiler_python.cpp。使用PyEval_SetProfile追蹤執行事件為了剖析調用棧實現使用PyEval_SetProfile來追蹤并處理 Python 程序中的各種執行事件對應源碼中recordPyCall/recordCCall的實現與PyTrace_*分支見 torch/csrc/autograd/profiler_python.cpp。它針對以下具體場景做出響應CPython 事件處理函數采集內容PyTrace_CALLrecordPyCall記錄每次 Python 函數調用捕獲后續分析所需的關鍵細節PyTrace_C_CALLrecordCCall記錄對 C 函數的調用包括相關參數提供程序執行流的完整視圖PyTrace_RETURN—記錄 Python 函數的退出時間實現精確的函數執行時長測量PyTrace_C_RETURN與PyTrace_C_EXCEPTION—記錄 C 函數的退出時間無論正常結束還是因異常結束確保所有執行路徑都被統計在RecordQueue/ThreadLocalSubqueue中Python 調用事件被以(TraceKey, approx_time_t)對的形式存入py_calls_torch/csrc/profiler/collection.h后處理階段通過匹配入口與出口來構建完整的調用事件。注意事項README 原文對于 Python 3.12.0–3.12.4CPython 中存在一個 bug需要使用sys.monitoring作為 workaround。這意味著不同 Python 小版本下 Python 棧追蹤的底層實現路徑可能不同。Python 前端的對應參數為with_stacktorch/profiler/profiler.pywith torch.profiler.profile( activities[torch.profiler.ProfilerActivity.CPU], with_stackTrue, # 記錄算子的源碼信息文件名與行號 with_modulesTrue, # 記錄模塊層級含函數名 ) as prof: ...需要說明的是README 與本倉庫當前代碼均提示with_modules已被標記為棄用deprecated未來的版本可能移除而舊式的with_flops參數在 eager 模式下已不起作用建議以with_stackTrue與record_shapesTrue組合來獲得可歸因的堆棧與形狀信息。時鐘對齊TSC 周期與 Unix 納秒之間的轉換在創建時間戳時Profiler 會根據系統環境選擇最有效的時鐘。大多數 Linux 系統的默認選擇是TSC它以 CPU 周期cycles的形式記錄時間。為了把這個時間轉換為納秒級 Unix 時間Profiler 會創建一個時鐘轉換器clock converter。如果 profiler 中包含了 Kineto這個轉換器也會被傳入 Kineto以確保兩端時間軸對齊——這正是 trace 中 CPU 事件與 GPU 事件能夠精確排布在統一時間線上的基礎。在源碼中可以看到這一機制的直接體現torch/csrc/autograd/profiler_kineto.cppauto converter clockConverter.makeConverter(); #ifdef USE_KINETO libkineto::get_time_converter() converter; #endif auto records_and_trace recordQueue.getRecords(std::move(converter), startTime, end_time);c10::ApproximateClockToUnixTimeConverterclockConverter的類型負責建立 TSC 計數與 Unix 時間的映射getRecords在把各線程子隊列中的原始近似時間事件物化為Result時統一使用該轉換器換算為納秒時間戳。相關類型定義位于 c10/util/ApproximateClock.h。總結一次完整 trace 的端到端數據流綜合以上各節一次torch.profiler.profile(...)調用的完整數據流可以歸納為啟動Python 前端profile()解析參數并調用后端profiler_kineto.cpp創建KinetoThreadLocalState與RecordQueue注冊onFunctionEnter/Exit回調全局或線程本地采集每個算子經 dispatcher 觸發RecordFunction回調 →begin_op在對應線程的ThreadLocalSubqueue中寫入輕量事件分配與 OOM 事件繞過RecordFunction直接入隊Python 棧事件由PyEval_SetProfile驅動GPU 事件由 Kineto 通過 CUPTI 等收集結束finalizeTrace()停止隊列、構建時鐘轉換器并同步給 KinetogetRecords()把原始事件物化為Result樹與 Kineto 事件合并、完成后處理前反向匹配、棧/模塊/形狀解碼、未完成事件補記結束時間導出通過ActivityTraceInterface::save寫出 JSON供 Chrome Trace / Perfetto 可視化或轉成FunctionEvents供 Python 側key_averages()等 API 做表格統計。理解這條鏈路后你將能更有針對性地使用 Profiler用record_shapes觀察算子輸入形狀、用with_stack定位調用來源、用profile_memory追蹤顯存分配并通過 Chrome Trace 中的 fwd_bwd flow 事件分析前反向傳播的時間占比。【免費下載鏈接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration項目地址: https://gitcode.com/GitHub_Trending/py/pytorch創作聲明：本文部分內容由AI輔助生成（AIGC），僅供參考 PyTorch Profiler 實現原理深度解析：從 RecordFunction 到 Kineto 的 CPU/GPU 全鏈路采集架構
PyTorch Profiler 實現原理深度解析從 RecordFunction 到 Kineto 的 CPU/GPU 全鏈路采集架構【免費下載鏈接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration項目地址: https://gitcode.com/GitHub_Trending/py/pytorchtorch.profiler是 PyTorch 官方提供的性能剖析工具它能夠記錄模型執行過程中的算子耗時、內存分配、Python 調用棧與 GPU 事件并導出為 Chrome Trace / Perfetto 可讀的 JSON 文件。本文以倉庫中 torch/csrc/profiler/README.md 為核心骨架結合 torch/csrc/profiler/collection.h、torch/csrc/profiler/collection.cpp、torch/csrc/autograd/profiler_kineto.cpp 等源碼深入剖析 Profiler 的底層實現事件從 CPU 算子回調、內存分配、自動求導、Python 棧到 GPU 活動是如何被采集、關聯與導出的。讀完本文你將理解 Profiler 的完整數據鏈路、各采集階段的內部機制以及如何基于這些原理正確配置和使用剖析參數。代碼庫結構Profiler 前端、后端與 C 采集層的職責劃分Profiler 的實現橫跨 Python 與 C 兩大層次README 給出了核心文件布局相對倉庫根路徑關鍵組成如下torch/ │ ├── profiler/ # 主 Python 包包含核心前端邏輯 │ ├── __init__.py # profiler 包的初始化文件 │ ├── profiler.py # 主 Profiler 前端類profile、schedule、supported_activities 等 │ └── _utils.py # FunctionEvent 工具函數 │ ├── autograd/ # Autograd 包 │ ├── __init__.py # autograd 包初始化文件 │ ├── profiler.py # 主 Profiler 后端類 │ └── profiler_utils.py # FunctionEvent 工具函數 │ ├── csrc/ # C 與 C 源碼 │ └── profiler/ # Profiler C 源碼 │ ├── collection.cpp # 主要采集邏輯 │ ├── collection.h # 采集定義 │ ├── kineto_client_interface.cpp # 從 kineto 調用 Profiler 的接口僅 on-demand 模式 │ ├── kineto_client_interface.h # Client 接口定義 │ ├── kineto_shim.cpp # 從 profiler 調用 kineto 的 shim │ ├── kineto_shim.h # shim 定義 │ ├── util.cpp # 處理 profiler 事件中參數的 utils │ ├── util.h # util 定義 │ └── README.md # 本文所講解的文檔 │ └── autograd/ # Autograd C 源碼 │ ├── profiler_python.cpp # 主要的 Python 調用棧采集邏輯 │ ├── profiler_python.h # Python 調用棧采集定義 │ ├── profiler_kineto.cpp # 啟動采集/kineto 的 Profiler 后端邏輯 │ └── profiler_kineto.h # 啟動采集/kineto 的 Profiler 后端定義 │ └── ATen/ # ATen C 源碼 │ ├── record_function.cpp # RecordFunction 采集邏輯 │ └── record_function.h # RecordFunction 定義從源碼結構可以清晰地看到職責邊界Python 前端torch/profiler/profiler.py負責向用戶暴露profile()、schedule()、ProfilerActivity等 API并把用戶的配置如record_shapes、with_stack翻譯成底層ProfilerConfig。Python 后端torch/autograd/profiler.pytorch.autograd.profiler模塊中的profile類歷史上是舊式 Profiler 的入口現在與 Kineto 后端并存。C 采集層torch/csrc/profiler/實現線程級事件緩沖ThreadLocalSubqueue、事件樹構建Result、Kineto shim 與時鐘轉換。Autograd C 層torch/csrc/autograd/profiler_kineto.cpp實現回調注冊與前后端橋接profiler_python.cpp實現基于 CPython C API 的調用棧追蹤。RecordFunctionCPU 側事件插樁的基礎設施RecordFunction是 Profiler 插樁 CPU 側事件的核心機制其定義位于 aten/src/ATen/record_function.h。它是一個通用的函數調用插樁手段并非 Profiler 專有也可以用于其他通用場景例如 PyTorch 的 大規模部署特性。在 PyTorch 內部它已被安插在若干關鍵位置最典型的是dispatcher中環繞每一個算子調用見 aten/src/ATen/core/dispatch/Dispatcher.h。RecordFunction的核心設計是回調注冊制用戶或 PyTorch 自身可以注冊回調每當執行路徑遇到一個RecordFunction守衛時這些回調就會被觸發。Profiler 正是利用這一機制來記錄每個算子調用的開始與結束時間以及用戶自定義的RecordFunction注釋。從工程角度RecordFunction機制被刻意設計為低開銷尤其是在沒有注冊任何回調時。這是因為算子調度路徑是訓練/推理的熱路徑任何額外開銷都會被放大。盡管如此一旦啟用回調依然會引入一定的性能損耗——這解釋了為什么 Profiler 只在需要時開啟采集。RecordFunction還提供了 Python 綁定with torch.profiler.record_function(...)。這是用戶在代碼中標記模塊級事件的常用手段例如import torch with torch.profiler.profile(activities[torch.profiler.ProfilerActivity.CPU]) as prof: with torch.profiler.record_function(my_custom_module): x torch.randn(1024, 1024) y x.mm(x)自定義的USER_SCOPE事件與算子事件一樣會被采集并出現在最終的 trace 中。Autograd 集成用序列號與前向線程 ID 關聯前反向自動求導引擎負責自動計算梯度。Profiler 從 autograd 引擎記錄兩類關鍵信息用于在 trace 中把反向算子與觸發它的前向算子關聯起來序列號Sequence Number定義位于 aten/src/ATen/SequenceNumber.h。這是一個每線程唯一的索引分配給前向傳播中的每次算子調用。當反向算子被觸發時它會被賦予與其來源前向算子相同的序列號。利用這一點Profiler 能夠完成前向/反向算子的匹配在 Chrome Trace 中該功能以fwd_bwd flow events的形式呈現。需要特別注意的是README 中的腳注只有輸入張量需要梯度的算子調用才會被分配序列號。這意味著純推理無梯度需求場景下序列號關聯的價值有限而在訓練場景中該機制是剖析反向傳播熱點的重要依據。前向線程 IDForward Thread IDAutograd 可以在多線程環境中使用。前向線程 ID 表示前向算子執行所在線程的 ID相關字段可見于 aten/src/ATen/record_function.h 的RecordFunction定義中。之所以需要它是因為上述序列號只在單個線程內唯一不同線程上可能產生相同的序列號必須用前向線程 ID 加以區分。在 torch/csrc/profiler/collection.h 的TorchOpBasicFields中可以看到這兩個字段的直接體現struct TorchOpBasicFields { int64_t sequence_number_{0}; uint64_t forward_tid_{0}; ... };ExtraFieldsEventType::TorchOp繼承TorchOpBasicFields在采集時隨事件一起寫入供后處理階段構建前反向關聯。Torch 算子采集流程回調注冊、線程子隊列與熱路徑優化本節描述auto-trace進程內、同步模式下 torch 算子的通用采集流程。關于 on-demand進程外、異步追蹤的細節可參考 Libkineto 的 READMEKineto 作為第三方子模塊被引入。回調的注冊profiler_kineto.cpp當一次 trace 開始時autograd/profiler 后端會調用 torch/csrc/autograd/profiler_kineto.cpp 來準備、啟動或停止采集。在 trace 啟動時定義于該文件中的onFunctionEnter與onFunctionExit回調會被注冊到RecordFunction機制中。源碼中回調分為兩組torch/csrc/autograd/profiler_kineto.cpponFunctionEnterGlobal/onFunctionExitGlobal全局回調用于KINETO_ONDEMAND或配置了profile_all_threads的KINETO會話onFunctionEnterTLS/onFunctionExitTLS線程本地回調僅對 trace 啟動時存在的線程生效。回調的注冊范圍由ExperimentalConfig決定分為兩種模式Global全局回調注冊到執行期間的所有線程。Local本地回調僅注冊到trace 開始時刻已存在的線程上。auto recordFunctionCallback at::RecordFunctionCallback(onFunctionEnterGlobal, onFunctionExitGlobal) .needsInputs(state_ptr-config().report_input_shapes) .scopes(scopes);needsInputs(config.report_input_shapes)表明只有當用戶開啟record_shapes時回調才需要算子輸入信息否則采集開銷更小。線程子隊列與begin_op在onFunctionEnter內部Profiler 會為每個線程創建一個ThreadLocalSubqueue實例見 torch/csrc/profiler/collection.h確保每個 CPU 算子都與它實際執行的線程關聯。當一個 torch 算子進入時Profiler 調用定義在collection.cpp的begin_op來記錄必要信息。ThreadLocalSubqueue的設計值得關注內部使用AppendOnlyList只追加、塊大小為 512作為事件存儲避免每次算子都分配獨立 vector每種事件類型對應獨立的存儲區torch_ops_算子事件、allocations_內存分配、backend_events_后端事件、vulkan_events_、ooms_OOM、py_calls_Python 調用算子事件的correlation id由EventBlock按塊起始 ID 塊內偏移計算無需額外分配torch/csrc/profiler/collection.cpp。README 特別強調begin_op被刻意設計得非常輕量因為它處于 profiling 的 hot path熱路徑上過高的額外開銷會扭曲 profile 結果、降低其參考價值。因此回調期間只收集最必要的信息大部分邏輯都推遲到后處理階段完成。這也是InputOutputEncoder存在的原因——它在采集期把算子輸入的形狀、dtype 編碼進連續的AppendOnlyList避免為每個算子創建 vector而在后處理時才解碼還原見 torch/csrc/profiler/collection.cpp。全局會話的并發安全對于全局回調profiler_kineto.cpp中還有一個GlobalCallbackSession協調機制torch/csrc/profiler/kineto_profiler.cpp由于全局回調會在任意線程上觸發而disableProfiler()可能在另一個線程上銷毀 profiler 狀態因此每個回調通過enter()/exit()只在極短的臨界區getGlobal()RecordQueue訪問內持有 in-flight 計數teardown 時drain()關閉會話并等待所有 in-flight 回調退出避免 use-after-free。退出回調還會校驗session_generation_丟棄跨會話懸空的退出事件。內存分配事件采集繞過 RecordFunction 的瞬時事件與有起止時刻的算子事件不同內存分配事件被表示為cpu_instant_event零持續時間。因此分配事件不經過RecordFunction而是由reportMemoryUsage直接調用emplace_allocation_event把事件入隊到對應的ThreadLocalSubqueue。在 torch/csrc/autograd/profiler_kineto.cpp 中可以看到void reportMemoryUsage( void* ptr, int64_t alloc_size, size_t total_allocated, size_t total_reserved, c10::Device device) override { if (config_.profile_memory !config_.disabled()) { recordQueue.getSubqueue()-emplace_allocation_event( c10::getApproximateTime(), ptr, alloc_size, total_allocated, total_reserved, device.type(), device.index()); } }同樣地OOM內存不足事件通過reportOutOfMemory→emplace_ooms_event進入隊列。RawAllocation結構torch/csrc/profiler/collection.h記錄了分配指針、大小、累計分配量/預留量以及設備信息并被static_assert(std::is_trivial_vRawAllocation)強制保持平凡類型以提升性能。在 Python 前端對應的開關是profile()的profile_memory參數torch/profiler/profiler.pytorch.profiler.profile( activities[torch.profiler.ProfilerActivity.CPU], profile_memoryTrue, # 跟蹤張量內存分配/釋放 )Kineto 集成多架構事件采集、預熱與導出Kineto 是 Profiler 獲取 GPU 及其他加速器事件的關鍵抽象層。它作為第三方子模塊third_party/kineto/目錄被引入通過與 CUPTI 等庫交互來接收 GPU 與加速器事件再轉發給前端 Profiler。預熱warmup的原因Kineto 需要時間prepare也稱 warmup這些第三方模塊以避免初始化例程扭曲 profile 結果。理論上可以在作業啟動時就完成預熱但讓 CUPTI 這類重量級庫持續運行會帶來顯著的不必要開銷因此預熱被安排在正式 trace 之前的短暫窗口內完成。雙向橋接與后處理如前所述profiler_kineto.cpp在后臺調用合適的 profiler 階段同時它還會調用kineto_shim.cpp后者觸發 Kineto 中的對應例程。當一次 trace 完成后Kineto 收集的所有事件會被轉發給 Profiler原因有兩個合并所有數據并完成 Profiler 事件與 Kineto 事件之間的后處理例如把 Kineto 事件嵌入Result事件樹見 torch/csrc/profiler/collection.h 的ExtraFieldsEventType::Kineto其中Flow結構用于將 Kineto 的 flow 事件映射進 profiler 樹把這些事件作為FunctionEvents轉發給 Python 前端供prof.key_averages()、prof.table()等 API 使用。文件導出集成的最后一步是文件導出。所有事件收集并后處理完成后它們可以導出為 JSON 文件供Perfetto或Chrome Tracer可視化。這一步通過調用 Kineto 的ActivityTraceInterface::save完成把所有事件信息寫入磁盤。在 Python 側tensorboard_trace_handler會自動生成 trace 目錄from torch.profiler import profile, ProfilerActivity, tensorboard_trace_handler with profile( activities[ProfilerActivity.CPU, ProfilerActivity.CUDA], on_trace_readytensorboard_trace_handler(./log/resnet18), ) as prof: model(x) # 輸出的 JSON 文件位于 ./log/resnet18/ 目錄下Python 棧追蹤基于 CPython 剖析 API 的實現當在 profiler 中設置with_stackTrue時Python 棧追蹤器會使用PythonTracerBase中定義的make函數生成具體實現位于profiler_python.cpptorch/csrc/autograd/profiler_python.cpp。使用PyEval_SetProfile追蹤執行事件為了剖析調用棧實現使用PyEval_SetProfile來追蹤并處理 Python 程序中的各種執行事件對應源碼中recordPyCall/recordCCall的實現與PyTrace_*分支見 torch/csrc/autograd/profiler_python.cpp。它針對以下具體場景做出響應CPython 事件處理函數采集內容PyTrace_CALLrecordPyCall記錄每次 Python 函數調用捕獲后續分析所需的關鍵細節PyTrace_C_CALLrecordCCall記錄對 C 函數的調用包括相關參數提供程序執行流的完整視圖PyTrace_RETURN—記錄 Python 函數的退出時間實現精確的函數執行時長測量PyTrace_C_RETURN與PyTrace_C_EXCEPTION—記錄 C 函數的退出時間無論正常結束還是因異常結束確保所有執行路徑都被統計在RecordQueue/ThreadLocalSubqueue中Python 調用事件被以(TraceKey, approx_time_t)對的形式存入py_calls_torch/csrc/profiler/collection.h后處理階段通過匹配入口與出口來構建完整的調用事件。注意事項README 原文對于 Python 3.12.0–3.12.4CPython 中存在一個 bug需要使用sys.monitoring作為 workaround。這意味著不同 Python 小版本下 Python 棧追蹤的底層實現路徑可能不同。Python 前端的對應參數為with_stacktorch/profiler/profiler.pywith torch.profiler.profile( activities[torch.profiler.ProfilerActivity.CPU], with_stackTrue, # 記錄算子的源碼信息文件名與行號 with_modulesTrue, # 記錄模塊層級含函數名 ) as prof: ...需要說明的是README 與本倉庫當前代碼均提示with_modules已被標記為棄用deprecated未來的版本可能移除而舊式的with_flops參數在 eager 模式下已不起作用建議以with_stackTrue與record_shapesTrue組合來獲得可歸因的堆棧與形狀信息。時鐘對齊TSC 周期與 Unix 納秒之間的轉換在創建時間戳時Profiler 會根據系統環境選擇最有效的時鐘。大多數 Linux 系統的默認選擇是TSC它以 CPU 周期cycles的形式記錄時間。為了把這個時間轉換為納秒級 Unix 時間Profiler 會創建一個時鐘轉換器clock converter。如果 profiler 中包含了 Kineto這個轉換器也會被傳入 Kineto以確保兩端時間軸對齊——這正是 trace 中 CPU 事件與 GPU 事件能夠精確排布在統一時間線上的基礎。在源碼中可以看到這一機制的直接體現torch/csrc/autograd/profiler_kineto.cppauto converter clockConverter.makeConverter(); #ifdef USE_KINETO libkineto::get_time_converter() converter; #endif auto records_and_trace recordQueue.getRecords(std::move(converter), startTime, end_time);c10::ApproximateClockToUnixTimeConverterclockConverter的類型負責建立 TSC 計數與 Unix 時間的映射getRecords在把各線程子隊列中的原始近似時間事件物化為Result時統一使用該轉換器換算為納秒時間戳。相關類型定義位于 c10/util/ApproximateClock.h。總結一次完整 trace 的端到端數據流綜合以上各節一次torch.profiler.profile(...)調用的完整數據流可以歸納為啟動Python 前端profile()解析參數并調用后端profiler_kineto.cpp創建KinetoThreadLocalState與RecordQueue注冊onFunctionEnter/Exit回調全局或線程本地采集每個算子經 dispatcher 觸發RecordFunction回調 →begin_op在對應線程的ThreadLocalSubqueue中寫入輕量事件分配與 OOM 事件繞過RecordFunction直接入隊Python 棧事件由PyEval_SetProfile驅動GPU 事件由 Kineto 通過 CUPTI 等收集結束finalizeTrace()停止隊列、構建時鐘轉換器并同步給 KinetogetRecords()把原始事件物化為Result樹與 Kineto 事件合并、完成后處理前反向匹配、棧/模塊/形狀解碼、未完成事件補記結束時間導出通過ActivityTraceInterface::save寫出 JSON供 Chrome Trace / Perfetto 可視化或轉成FunctionEvents供 Python 側key_averages()等 API 做表格統計。理解這條鏈路后你將能更有針對性地使用 Profiler用record_shapes觀察算子輸入形狀、用with_stack定位調用來源、用profile_memory追蹤顯存分配并通過 Chrome Trace 中的 fwd_bwd flow 事件分析前反向傳播的時間占比。【免費下載鏈接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration項目地址: https://gitcode.com/GitHub_Trending/py/pytorch創作聲明：本文部分內容由AI輔助生成（AIGC），僅供參考