Lightning-AI / Lightning-AI/pytorch-lightning

Callback for logging forward, backward and update time

Open
#19,939 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

callback feature profiler
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Description & Motivation

Cross-posting an interesting and important feature request from the discussions channel: https://github.com/Lightning-AI/pytorch-lightning/discussions/19928

Citing the discussion thread:
> I'm trying to track the performance of forward/backward/update time with a Callback. My current implementation is showing strange behavior. It seems that the callback order is (at least functionally) different when using gradient accumulation !=1. This is expected, but it's unclear how to handle both cases with a single callback.

### Pitch

Citing the discussion thread:
> I'd really appreciate help coming up with an almost-correct implementation for tracking (1) forward pass time, (2) backward pass time, (3) total time for an update (forward + backward + optimizer step) which might depend on gradient accumulation and (4) amount of time spent waiting on the dataloader to generate the next batch. Alternatively, for (3) I'm happy to track only optimizer.step time since this should tell me how long it's taking for devices to sync and gradients to update. I'm open to tracking related metrics or other metrics entirely as long as they're correlated with model throughput/performance.

### Alternatives

An implementation is proposed on the discussion thread, but the OP shows that it is in fact incorrect.

### Additional context

Please check out the discussion! https://github.com/Lightning-AI/pytorch-lightning/discussions/19928

cc @lantiga @borda

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with discussion #19928 and the implementation proposed there, then compare callback behavior with and without gradient accumulation. Done means agreeing on a callback design that measures forward, backward, update or optimizer-step, and dataloader waiting time correctly across accumulation settings.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.