Repository: Lightning-AI/pytorch-lightning
Pull request: Improve callback documentation for `outputs` and `accumulate_grad_batches` (Resolves #15315)
## What does this PR do?

Adds brief documentation explaining the `outputs` parameter to callback hooks, and that potentially unexpected behavior occurs when using `accumulate_grad_batches` > 1 

Fixes #15315 

### Does your PR introduce any breaking changes? If yes, please list them.

None

## Before submitting

- [X] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [X] Did you read the [contributor guideline](https://github.com/Lightning-AI/lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [X] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- [X] Did you make sure to **update the documentation** with your changes? (if necessary)
- [ ] Did you write any **new necessary tests**? (not for typos and docs)
- [X] Did you verify new and **existing tests pass** locally with your changes?
- [X] Did you list all the **breaking changes** introduced by this pull request?
- [X] Did you **update the CHANGELOG**? (not for typos, docs, test updates, or minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [ ] Is this pull request ready for review? (if not, please submit in draft mode)
- [ ] Check that all items from **Before submitting** are resolved
- [ ] Make sure the title is self-explanatory and the description concisely explains the PR
- [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?

Make sure you had fun coding 🙃


Problem statement:
Misleading outputs (loss) in callback when using `accumulate_grad_batches` > 1 in trainer
## 📚 Documentation

Some of the callback entry points/hooks (e.g. `on_train_batch_end`) have an `outputs` argument that contains the loss returned from the `training_step` and `validation_step`. This value does not match the loss in the trainer when `accumulate_grad_batches` > 1. I spoke with @rohitgr7 about this on slack and he said that this is expected behavior and that the loss is normalized internally. We agreed the best fix was a documentation PR alerting users of this possibility, and creating an issue so people could find it if they also experience it. I will submit a PR updating the documentation for the callback entry points explaining this behavior.

Hints from issue discussion:

Commit message: Update src/pytorch_lightning/core/module.py

Co-authored-by: Rohit Gupta <rohitgr1998@gmail.com>

Patch to review:
diff --git a/docs/source-pytorch/extensions/callbacks.rst b/docs/source-pytorch/extensions/callbacks.rst
index 401b9606258b6..106e718bb1181 100644
--- a/docs/source-pytorch/extensions/callbacks.rst
+++ b/docs/source-pytorch/extensions/callbacks.rst
@@ -138,6 +138,11 @@ Callback API
 ************
 Here is the full API of methods available in the Callback base class.
 
+.. note::
+
+    Outputs: Hooks that occur on batch end (e.g. ``on_train_batch_end``) will receive any values returned from that step (such as the loss) in the form of a dictionary named ``outputs``.
+    ``outputs["loss"]`` here will be the normalized value w.r.t ``accumulate_grad_batches`` of the loss returned from ``training_step``.
+
 The :class:`~pytorch_lightning.callbacks.Callback` class is the base for all the callbacks in Lightning just like the :class:`~pytorch_lightning.core.module.LightningModule` is the base for all models.
 It defines a public interface that each callback implementation must follow, the key ones are:
 
diff --git a/src/pytorch_lightning/core/module.py b/src/pytorch_lightning/core/module.py
index c557d8fb0f623..7710e4d5c6b86 100644
--- a/src/pytorch_lightning/core/module.py
+++ b/src/pytorch_lightning/core/module.py
@@ -661,6 +661,10 @@ def training_step(self, batch, batch_idx, hiddens):
         Note:
             The loss value shown in the progress bar is smoothed (averaged) over the last values,
             so it differs from the actual loss returned in train/validation step.
+
+        Note:
+            When ``accumulate_grad_batches`` > 1, the loss returned here will be automatically
+            normalized by ``accumulate_grad_batches`` internally.
         """
         rank_zero_warn("`training_step` must be implemented to be used with the Lightning Trainer")
 


Code review comments:
diff_hunk: @@ -138,6 +138,11 @@ Callback API
 ************
 Here is the full API of methods available in the Callback base class.
 
+.. note::
+
+    Outputs: Hooks that occur on batch end (e.g. ``on_train_batch_end``) will receive any values returned from that step (such as the loss) in the form of a dictionary named ``outputs``.
+    ``outputs["loss"]`` here will be the normalized value w.r.t ``accumulate_grad_batches`` of the loss returned from ``training_step``.
path: docs/source-pytorch/extensions/callbacks.rst
text: @user1:
This is not the right place to put it IMO. 

1. This note is inserted before we explain the Callback: "The Callback class is the base for all the callbacks in Lightning". This seems distracting.
2. I propose to add a simple sentence explaining this in the individual hook documentation. When a user implements the  on_train_batch_end, they would go to the documentation of that hook if they have any question about its signature. 

@author:
@user1, thanks for the suggestion, this change has been pushed and is ready for review.