PyTorchProfiler: not showing CPU memory used even with `profile_memory=True`

[Jack12xl]

Bug description

Trying to use PyTorchProfiler (https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.profilers.PyTorchProfiler.html) to track some OOM(cpu memory) issues.
I go with

profiler = PyTorchProfiler(
        dirpath=log_dir,  # Directory to save logs
        filename="memory_profile",  # Name of the file to save results
        sort_by_key="self_cpu_memory_usage",  # Sort by CPU memory usage
        export_to_chrome=True,  # Export as JSON for Chrome
        row_limit=16,
        activities=[torch.profiler.ProfilerActivity.CPU],
        profile_memory=True,  # Record CPU memory usage
        with_stack=True,
        record_shapes=True,
    )

trainer=pl.Trainer(..., profiler=profiler, ...)
trainer.fit()

I expected to results similar to Pytorch Native profiler(https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html#using-profiler-to-analyze-memory-consumption). But it's still outputting cpu/gpu time like here:

I don't know if this is a bug(I thought PyTorchProfiler was a wrapper around native Pytorch Profiler, so it should have similar behavior when I set profile_memory=True).

Thanks! Please correct me if I am wrong!

What version are you seeing the problem on?

v2.4

How to reproduce the bug

profiler = PyTorchProfiler(
dirpath=log_dir, # Directory to save logs
filename="memory_profile", # Name of the file to save results
sort_by_key="self_cpu_memory_usage", # Sort by CPU memory usage
export_to_chrome=True, # Export as JSON for Chrome
row_limit=16,
activities=[torch.profiler.ProfilerActivity.CPU],
profile_memory=True, # Record CPU memory usage
with_stack=True,
record_shapes=True,
)

trainer=pl.Trainer(..., profiler=profiler, ...)
trainer.fit()

Error messages and logs

# Error messages and logs here please

Environment

Current environment

* CUDA:
	- GPU:
		- NVIDIA 30xx GPU
	- available:         True
	- version:           12.1
* Lightning:
	- lightning:         2.4.0
	- lightning-utilities: 0.11.7
	- pytorch-lightning: 2.4.0
	- torch:             2.3.1
	- torchaudio:        2.3.1
	- torchdata:         0.8.0
	- torchmetrics:      1.4.1
	- torchvision:       0.18.1

Python: 3.12.4

More info

It's not directly related to this issue. But is there some way I could have the export_memory_timeline(https://pytorch.org/docs/main/profiler.html#torch.profiler._KinetoProfile.export_memory_timeline) behavior with lightning PytorchProfiler?

Thanks!

[LTMeyer]

It's not directly related to this issue. But is there some way I could have the export_memory_timeline(https://pytorch.org/docs/main/profiler.html#torch.profiler._KinetoProfile.export_memory_timeline) behavior with lightning PytorchProfiler?

You can be using the on_trace_ready kwargs of the lightning.pytorch.profilers.PyTorchProfiler.

from torch.profiler import tensorboard_trace_handler, profile
from lightning.pytorch.profilers import PyTorchProfiler


def memory_trace_handler(dir_name: str):
    """Handler to export memory timeline to HTML.

    Args:
        dir_name: Directory to save the trace.
    """
    file_name = f"{dir_name}/memory_timeline.html"

    def handler_fn(prof: profile):
        prof.export_memory_timeline(file_name)

    return handler_fn


profiler = PyTorchProfiler(profile_memory=True, on_trace_ready=memory_trace_handler(dir_name))
# Use profiler with trainer
...

[SkafteNicki]

Thanks for raising this issue. I am unsure when this was fixed but I can happily report that if you try to use the PytorchProfiler in the latest version of lightning, this now correctly shows the CPU memory usage. An example script:

import os
import tempfile
import torch
import torch.nn as nn
import lightning as L
from lightning.pytorch.profilers import PyTorchProfiler
from lightning.pytorch import Trainer
from lightning.pytorch.demos.boring_classes import BoringModel


def main():
    # Create temporary directory for logs
    with tempfile.TemporaryDirectory() as log_dir:
        
        # Exact configuration from the issue
        profiler = PyTorchProfiler(
            dirpath=log_dir,  # Directory to save logs
            filename="memory_profile",  # Name of the file to save results
            sort_by_key="self_cpu_memory_usage",  # Sort by CPU memory usage
            export_to_chrome=True,  # Export as JSON for Chrome
            row_limit=16,
            activities=[torch.profiler.ProfilerActivity.CPU],
            profile_memory=True,  # Record CPU memory usage
            with_stack=True,
            record_shapes=True,
        )
        
        # Simple model to test with
        class TestModel(BoringModel):
            def training_step(self, batch, batch_idx):
                x = batch
                # Create memory pressure similar to OOM scenarios
                large_tensor = torch.randn(800, 800, device=x.device)  # ~2.5MB
                y_hat = self.layer(x) 
                intermediate = y_hat + large_tensor.sum() * 1e-6  # Small contribution
                del large_tensor
                loss = nn.functional.mse_loss(intermediate, torch.ones_like(intermediate))
                return loss
        
        model = TestModel()
        trainer = Trainer(
            max_epochs=1,
            limit_train_batches=10,  # More batches to get better profiling
            limit_val_batches=5,
            profiler=profiler,
            logger=False,
            enable_checkpointing=False,
            enable_progress_bar=False,
            accelerator="cpu"
        )
        
        trainer.fit(model)
        
        print("\nProfiler output:")
        print(profiler.summary())
        
        # Check files created
        print(f"\nFiles created in {log_dir}:")
        files = os.listdir(log_dir)
        for file in files:
            print(f"  - {file}")
            if file.endswith('.json'):
                print(f"    ✅ Chrome trace file created: {file}")
        
        # Analyze the output
        summary = profiler.summary()
        has_memory_columns = "CPU Mem" in summary or "Self CPU Mem" in summary
        is_sorted_by_memory = "self_cpu_memory_usage" in str(profiler._sort_by_key)
        
        print(f"\nAnalysis:")
        print(f"  Memory columns present: {has_memory_columns}")
        print(f"  Sort key is memory-related: {is_sorted_by_memory}")
        print(f"  Expected sort key: self_cpu_memory_usage")
        print(f"  Actual sort key: {profiler._sort_by_key}")
        
        return has_memory_columns


if __name__ == "__main__":
    main()

which when run will show something like this:

Profiler output:
FIT Profiler Report
Profile stats for: records
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                              aten::cat         0.69%      23.363us         0.69%      23.363us      11.682us         256 B         256 B             2  
                                    aten::empty_strided         0.44%      14.990us         0.44%      14.990us       3.747us          24 B          24 B             4  
                                            aten::addmm         1.27%      43.065us         2.12%      71.722us      35.861us          16 B          16 B             2  
                                            aten::empty         0.17%       5.819us         0.17%       5.819us       2.909us          16 B          16 B             2  
                                             aten::div_         0.64%      21.606us         1.27%      43.006us      21.503us          16 B           8 B             2  
[pl][profile][Callback]RichModelSummary.on_validatio...         3.61%     122.112us         3.61%     122.112us      61.056us           0 B           0 B             2  
[pl][profile][LightningModule]TestModel.on_validatio...         1.93%      65.389us         1.93%      65.389us      32.695us           0 B           0 B             2  
                [pl][profile][_EvaluationLoop].val_next         3.46%     116.885us        11.84%     400.306us     200.153us         256 B           0 B             2  
enumerate(DataLoader)#_SingleProcessDataLoaderIter._...         5.91%     199.911us         8.38%     283.421us     141.711us         256 B           0 B             2  
                                           aten::select         0.93%      31.416us         1.09%      36.929us      18.465us           0 B           0 B             2  
                                       aten::as_strided         0.35%      11.779us         0.35%      11.779us       1.472us           0 B           0 B             8  
                                            aten::stack         0.52%      17.591us         1.38%      46.581us      23.291us         256 B           0 B             2  
                                             aten::view         0.17%       5.627us         0.17%       5.627us       2.813us           0 B           0 B             2  
[pl][profile][LightningModule]TestModel.on_before_ba...         1.78%      60.334us         1.78%      60.334us      30.167us           0 B           0 B             2  
[pl][profile][Strategy]SingleDeviceStrategy.batch_to...         5.58%     188.669us        10.32%     348.893us     174.447us           0 B           0 B             2  
[pl][profile][LightningModule]TestModel.transfer_bat...         3.09%     104.582us         3.15%     106.471us      53.236us           0 B           0 B             2  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------

with the CPU Mem Self CPU Mem columns included. Closing issue for now, but feel free to ping and reopen if necessary.