Assertion failed in torch.profiler.export_memory_timeline()

Conless: <p>Hi! I was using torch.profiler.profile to profile the memory usage of my training code, which consumes more memory than expected. So I use the profiler to wrap my training code as what is done in the example:</p>
<pre data-code-wrap="python"><code class="lang-python">    def trace_handler(prof: torch.profiler.profile):

            # Prefix for file names.
            host_name = socket.gethostname()
            TIME_FORMAT_STR: str = "%b_%d_%H_%M_%S"
            timestamp = datetime.now().strftime(TIME_FORMAT_STR)
            file_prefix = f"{host_name}_{timestamp}"

            # Construct the trace file.
            prof.export_chrome_trace(f"trace/{file_prefix}.json.gz")

            # Construct the memory timeline file.
            prof.export_memory_timeline(f"trace/{file_prefix}.html", device="cuda:0")
        
        with torch.profiler.profile(
            activities=[
                torch.profiler.ProfilerActivity.CPU,
                torch.profiler.ProfilerActivity.CUDA,
            ],
            schedule=torch.profiler.schedule(wait=0, warmup=0, active=6, repeat=1),
            record_shapes=True,
            profile_memory=True,
            with_stack=True,
            on_trace_ready=trace_handler,
        ) as prof:
            tr_losses = trainer.train()
</code></pre>
<p>However, when I run the code by <code>torchrun</code>, it fails with traceback:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/app/train.py", line 371, in test_forward_with_adapter_hook
    with torch.profiler.profile(
  File "/opt/conda/envs/finetune/lib/python3.10/site-packages/torch/profiler/profiler.py", line 706, in __exit__
    self.stop()
  File "/opt/conda/envs/finetune/lib/python3.10/site-packages/torch/profiler/profiler.py", line 722, in stop
    self._transit_action(self.current_action, None)
  File "/opt/conda/envs/finetune/lib/python3.10/site-packages/torch/profiler/profiler.py", line 751, in _transit_action
    action()
  File "/opt/conda/envs/finetune/lib/python3.10/site-packages/torch/profiler/profiler.py", line 745, in _trace_ready
    self.on_trace_ready(self)
  File "/app/_test_adapter_hook_on_sharded_model.py", line 369, in trace_handler
    prof.export_memory_timeline(f"trace/{file_prefix}.html", device="cuda:0")
  File "/opt/conda/envs/finetune/lib/python3.10/site-packages/torch/profiler/profiler.py", line 338, in export_memory_timeline
    self.mem_tl = MemoryProfileTimeline(self._memory_profile())
  File "/opt/conda/envs/finetune/lib/python3.10/site-packages/torch/profiler/profiler.py", line 306, in _memory_profile
    return MemoryProfile(self.profiler.kineto_results)
  File "/opt/conda/envs/finetune/lib/python3.10/site-packages/torch/profiler/_memory_profiler.py", line 665, in __init__
    self._data_flow_graph = DataFlowGraph(self._op_tree)
  File "/opt/conda/envs/finetune/lib/python3.10/site-packages/torch/profiler/_memory_profiler.py", line 518, in __init__
    self._flow_nodes = [DataFlowNode(e, self) for e in self.leaf_events]
  File "/opt/conda/envs/finetune/lib/python3.10/site-packages/torch/profiler/_memory_profiler.py", line 518, in &lt;listcomp&gt;
    self._flow_nodes = [DataFlowNode(e, self) for e in self.leaf_events]
  File "/opt/conda/envs/finetune/lib/python3.10/site-packages/torch/profiler/_memory_profiler.py", line 431, in __init__
    assert all(i == j for i, j in versions.values()), f"{versions}, {self._edges}"
AssertionError: {id=6716:         0x3e093980 (7843) (cpu): (0, 2)}, {id=6716:         0x3e093980 (7843) (cpu): DataFlowEdge(input_version=None, mutated=False), id=8076:         0x3e739b90 (9445) (cpu): DataFlowEdge(input_version=0, mutated=False)}
</code></pre>
<p>What may be the problem? Thanks a lot!</p>