Skip to content

docs: add aggregate profiler memory example #4054

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
SunMarc merged 1 commit into from
Jun 9, 2026
+34 −1
Merged

docs: add aggregate profiler memory example #4054

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 34 additions & 1 deletion docs/source/usage_guides/profiler.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,39 @@ The resulting table output (omitting some columns):
Self CPU time total: 69.332ms
```

If you need one or two aggregate numbers for experiment tracking instead of a full operator table, combine the profiler with the runtime memory stats exposed by PyTorch. This is often enough for logging peak memory usage per run:

```python
import torch
from accelerate import Accelerator, ProfileKwargs


def summarize_memory_stats():
return {
"peak_allocated_mb": torch.cuda.max_memory_allocated() / 1024**2,
"peak_reserved_mb": torch.cuda.max_memory_reserved() / 1024**2,
}


profile_kwargs = ProfileKwargs(
activities=["cpu", "cuda"],
profile_memory=True,
record_shapes=True,
)

accelerator = Accelerator(kwargs_handlers=[profile_kwargs])

torch.cuda.reset_peak_memory_stats()

with accelerator.profile() as prof:
model(inputs)

print(prof.key_averages().table(sort_by="self_cuda_memory_usage", row_limit=10))
print(summarize_memory_stats())
```

This pattern keeps the profiler output available for debugging while also producing compact values that can be logged to an experiment tracker.


## Exporting chrome trace

Expand Down Expand Up @@ -334,4 +367,4 @@ Self CUDA time total: 4.165ms

PyTorch Profiler is a powerful tool for analyzing the performance of your models. By integrating it with Accelerate, you can easily profile your models and gain insights into their performance, helping you to optimize and improve them.

For more detailed information, refer to the [PyTorch Profiler documentation](https://pytorch.org/docs/stable/profiler.html).
For more detailed information, refer to the [PyTorch Profiler documentation](https://pytorch.org/docs/stable/profiler.html).
Loading