[Profiler] Account for caching when assigning IDs (#88917)

Taylor Robie · pytorchmergebot · commit 1db0f735e8fe · 2022-11-15T18:24:15.000Z
The python tracer caches information about module and optimizer state. That means that for subsequent calls, the presence of a Tensor in these fields does not imply that the Tensor is still live; just that it was live during the first call. (I should perhaps rename the fields to something like `stale_parameters` to convey this.) Unless we discard subsequent calls ID assignment get tripped up when it see's a Tensor that was already released. Differential Revision: [D41226827](https://our.internmc.facebook.com/intern/diff/D41226827/) Pull Request resolved: #88917 Approved by: https://github.com/chaekit
diff --git a/torch/csrc/profiler/data_flow.cpp b/torch/csrc/profiler/data_flow.cpp
@@ -69,6 +69,10 @@ void calculateUniqueTensorIDs(
   // --------------------------------------------------------------------------
   {
     RawTensors raw_tensors;
+
+    // The python tracer caches values, so it's only safe to use the first case.
+    ska::flat_hash_set<PyModuleSelf> seen_modules;
+    ska::flat_hash_set<PyOptimizerSelf> seen_optimizers;
     for (auto& result : sorted_results) {
       result->visit(c10::overloaded(
           [&](ExtraFields<EventType::TorchOp>& torch_op) {
@@ -78,15 +82,17 @@ void calculateUniqueTensorIDs(
           },
           [&](ExtraFields<EventType::PyCall>& py_call) {
             // torch.nn.Module
-            if (py_call.module_.has_value()) {
+            if (py_call.module_.has_value() &&
+                seen_modules.insert(py_call.module_->self_).second) {
               for (auto& p : py_call.module_->parameters_) {
                 raw_tensors(p.metadata_);
                 raw_tensors(p.grad_metadata_);
               }
             }
 
             // torch.optim.Optimizer
-            if (py_call.optimizer_.has_value()) {
+            if (py_call.optimizer_.has_value() &&
+                seen_optimizers.insert(py_call.optimizer_->self_).second) {
               for (auto& p : py_call.optimizer_->parameters_) {
                 raw_tensors(p.metadata_);
                 raw_tensors(p.grad_metadata_);
