Improve docs for Callbacks on access to all batch outputs

[dudarboh]

📚 Documentation

Hi Lightning team,

The docs encourage that the Callbacks should be "NON-ESSENTIAL logic that is NOT required for your lightning module".

However, in the example of accessing all batch outputs in on_train_epoch_end() docs, the clearing of the memory is done inside the callback, making it essential to avoid memory build-up.

It would be nice to adapt this example so that memory clean-up does not depend on the Callback.

I would assume it is as easy as moving memory-clearing code to the LightningModule's on_train_epoch_end(), however, I couldn't find any documentation of what is being called first, the LightningModule's on_train_epoch_end() or the Callbacks... Maybe I have missed it, but if not, then it would be another improvement to the docs.

Cheers

cc @lantiga @Borda

[bhimrazy]

Agreed, @dudarboh — having a sequence diagram for the event flow would make this much easier to understand.

Regarding the order of execution, based on the code here:

pytorch-lightning/src/lightning/pytorch/loops/fit_loop.py

Lines 463 to 471 in a9deada

	# call train epoch end hooks
	# we always call callback hooks first, but here we need to make an exception for the callbacks that
	# monitor a metric, otherwise they wouldn't be able to monitor a key logged in
	# `LightningModule.on_train_epoch_end`
	call._call_callback_hooks(trainer, "on_train_epoch_end", monitoring_callbacks=False)
	call._call_lightning_module_hook(trainer, "on_train_epoch_end")
	call._call_callback_hooks(trainer, "on_train_epoch_end", monitoring_callbacks=True)

---

So, the cleanup logic could also be placed directly in the LightningModule if the hook is defined there , without depending on callback.

Example: https://lightning.ai/docs/pytorch/stable/extensions/callbacks.html#on-train-epoch-end

class MyLightningModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.training_step_outputs = []

    def training_step(self):
        loss = ...
        self.training_step_outputs.append(loss)
        return loss

    def on_train_epoch_end(self):
        # aggregate outputs
        epoch_mean = torch.stack(self.training_step_outputs).mean()
        self.log("training_epoch_mean", epoch_mean)
        # clear to free memory
        self.training_step_outputs.clear()

---

In the other callback example, the cleanup is shown in a callback, which likely leaves it to the user to decide whether to handle it in the LightningModule or in the callback. But yes — having the precise event order documented would help make that decision with confidence.

[SkafteNicki]

@bhimrazy @dudarboh would this be usefull to have in docs:

    Training Process Flow:
    
    trainer.fit()
    │
    ├── setup(stage="fit")
    │   └── [Callbacks only]
    │
    ├── on_fit_start()
    │   ├── [Callbacks]
    │   ├── [LightningModule]  
    │   └── [Strategy]
    │
    ├── on_sanity_check_start()
    │   ├── [Callbacks]
    │   ├── [LightningModule]
    │   └── [Strategy]
    │   ├── on_validation_start()
    │   │   ├── [Callbacks]
    │   │   ├── [LightningModule]
    │   │   └── [Strategy]
    │   ├── on_validation_epoch_start()
    │   │   ├── [Callbacks]
    │   │   ├── [LightningModule]
    │   │   └── [Strategy]
    │   │   ├── [for each validation batch]
    │   │   │   ├── on_validation_batch_start()
    │   │   │   │   ├── [Callbacks]
    │   │   │   │   ├── [LightningModule]
    │   │   │   │   └── [Strategy]
    │   │   │   └── on_validation_batch_end()
    │   │   │       ├── [Callbacks]
    │   │   │       ├── [LightningModule]
    │   │   │       └── [Strategy]
    │   │   └── [end validation batches]
    │   ├── on_validation_epoch_end()
    │   │   ├── [Callbacks]
    │   │   ├── [LightningModule]
    │   │   └── [Strategy]
    │   └── on_validation_end()
    │       ├── [Callbacks]
    │       ├── [LightningModule]
    │       └── [Strategy]
    ├── on_sanity_check_end()
    │   ├── [Callbacks]
    │   ├── [LightningModule]
    │   └── [Strategy]
    │
    ├── on_train_start()
    │   ├── [Callbacks]
    │   ├── [LightningModule]
    │   └── [Strategy]
    │
    ├── [Training Epochs Loop]
    │   │
    │   ├── on_train_epoch_start()
    │   │   ├── [Callbacks]
    │   │   └── [LightningModule]
    │   │
    │   ├── [Training Batches Loop]
    │   │   │
    │   │   ├── on_train_batch_start()
    │   │   │   ├── [Callbacks]
    │   │   │   ├── [LightningModule]
    │   │   │   └── [Strategy]
    │   │   │
    │   │   ├── on_before_zero_grad()
    │   │   │   ├── [Callbacks]
    │   │   │   └── [LightningModule]
    │   │   │
    │   │   ├── [Forward Pass - training_step()]
    │   │   │   └── [Strategy only]
    │   │   │
    │   │   ├── on_before_backward()
    │   │   │   ├── [Callbacks]
    │   │   │   └── [LightningModule]
    │   │   │
    │   │   ├── [Backward Pass]
    │   │   │   └── [Strategy only]
    │   │   │
    │   │   ├── on_after_backward()
    │   │   │   ├── [Callbacks]
    │   │   │   └── [LightningModule]
    │   │   │
    │   │   ├── on_before_optimizer_step()
    │   │   │   ├── [Callbacks]
    │   │   │   └── [LightningModule]
    │   │   │
    │   │   ├── [Optimizer Step]
    │   │   │   └── [LightningModule only - optimizer_step()]
    │   │   │
    │   │   └── on_train_batch_end()
    │   │       ├── [Callbacks]
    │   │       └── [LightningModule]
    │   │   
    │   │   [Optional: Validation during training]
    │   │   ├── on_validation_start()
    │   │   │   ├── [Callbacks]
    │   │   │   ├── [LightningModule]
    │   │   │   └── [Strategy]
    │   │   ├── on_validation_epoch_start()
    │   │   │   ├── [Callbacks]
    │   │   │   ├── [LightningModule]
    │   │   │   └── [Strategy]
    │   │   │   ├── [for each validation batch]
    │   │   │   │   ├── on_validation_batch_start()
    │   │   │   │   │   ├── [Callbacks]
    │   │   │   │   │   ├── [LightningModule]
    │   │   │   │   │   └── [Strategy]
    │   │   │   │   └── on_validation_batch_end()
    │   │   │   │       ├── [Callbacks]
    │   │   │   │       ├── [LightningModule]
    │   │   │   │       └── [Strategy]
    │   │   │   └── [end validation batches]
    │   │   ├── on_validation_epoch_end()
    │   │   │   ├── [Callbacks]
    │   │   │   ├── [LightningModule]
    │   │   │   └── [Strategy]
    │   │   └── on_validation_end()
    │   │       ├── [Callbacks]
    │   │       ├── [LightningModule]
    │   │       └── [Strategy]
    │   │
    │   └── on_train_epoch_end() **SPECIAL CASE**
    │       ├── [Callbacks - Non-monitoring only]
    │       ├── [LightningModule]
    │       └── [Callbacks - Monitoring only]
    │
    ├── [End Training Epochs]
    │
    ├── on_train_end()
    │   ├── [Callbacks]
    │   ├── [LightningModule]
    │   └── [Strategy]
    │
    ├── on_fit_end()
    │   ├── [Callbacks]
    │   ├── [LightningModule]
    │   └── [Strategy]
    │
    └── teardown(stage="fit")
        └── [Callbacks only]

[Borda]

would this be usefull to have in docs

Did you draw it manually or use an automated tracing tool?

[SkafteNicki]

would this be usefull to have in docs

Did you draw it manually or use an automated tracing tool?

I use an AI agent to track calls throughout the codebase and asked it to return the result as a simple call graph. Should be checked for correctness of cause before included in docs. Similar graph for trainer.test

    trainer.test()
    │
    ├── setup(stage="test")
    │   └── [Callbacks only]
    ├── on_test_start()
    │   ├── [Callbacks]
    │   ├── [LightningModule]
    │   └── [Strategy]
    │
    ├── [Test Epochs Loop]
    │   │
    │   ├── on_test_epoch_start()
    │   │   ├── [Callbacks]
    │   │   ├── [LightningModule]
    │   │   └── [Strategy]
    │   │
    │   ├── [Test Batches Loop]
    │   │   │
    │   │   ├── on_test_batch_start()
    │   │   │   ├── [Callbacks]
    │   │   │   ├── [LightningModule]
    │   │   │   └── [Strategy]
    │   │   │
    │   │   └── on_test_batch_end()
    │   │       ├── [Callbacks]
    │   │       ├── [LightningModule]
    │   │       └── [Strategy]
    │   │
    │   └── on_test_epoch_end()
    │       ├── [Callbacks]
    │       ├── [LightningModule]
    │       └── [Strategy]
    │
    ├── on_test_end()
    │   ├── [Callbacks]
    │   ├── [LightningModule]
    │   └── [Strategy]
    └── teardown(stage="test")
        └── [Callbacks only]

[bhimrazy]

Yes @SkafteNicki, I think this could be really helpful for understanding the flow.

There’s also another reference here from deepwiki generated version: Callback Lifecycle and Hook Points

cc: @Borda

[SkafteNicki]

There’s also another reference here from deepwiki generated version: Callback Lifecycle and Hook Points

Alright then I will be creating a PR with a new page on hook call order over the lifetime of training/testing/prediction.
And @bhimrazy thanks for linking to deepwiki, was not aware of that page.