[RFC] Deprecate `weights_summary` off the Trainer constructor

[ananthsub]

Proposed refactoring or deprecation

• Introduce a new ModelSummary callback, which calls summarize: https://github.com/PyTorchLightning/pytorch-lightning/blob/8a931732ae5135e3e55d9c7b7031d81837e5798a/pytorch_lightning/utilities/model_summary.py#L437-L439
• Deprecate weights_summary off the Trainer constructor

Motivation

We are auditing the Lightning components and APIs to assess opportunities for improvements:

• Review Lightning architecture & API #7740
• https://docs.google.com/document/d/1xHU7-iQSpp9KJTjI3As2EM0mfNHHr37WZYpDpwLkivA/edit#

This is a followup to #8478 and #9006

Why do we want to remove this from the core trainer logic?

• We need a way for users to customize more of the inputs to the model summary over time without affecting the trainer API. Today, changes to the model summary API also require changes in the core trainer (e.g. the addition of max_depth ). This gives model summarization more room to grow without cascading changes elsewhere.
• Users may want to configure this summarization for different points of execution. Right now, this is hardcoded to be run only during fit(). But users could want to call this potentially multiple times during each of trainer.fit(), trainer.validate(), trainer.test() or trainer.predict().
• Users may want to customize where they save the summary. Right now, it's printed to stdout, but this could also be useful to save to a file or upload to another service for tracking the run.
• The current implementation runs on global rank 0 only in order to avoid printing out multiple summary tables. However, running this on rank 0 will break for model parallel use cases that require communication across ranks. This can lead to subtle failures if example_input_array is set as a property on the LightningModule. For instance, a model wrapped with FSDP will break because parameters need to be all-gathered across layers across ranks.
• In case the LightningModule leverages PyTorch LazyModules, users may want to generate this summary only after the first batch is processed in order to get accurate parameter estimations. Estimates of parameter sizes with lazy modules would be misleading.
• AFAICT, this is the only piece of logic that runs in between on_pretrain_routine_start/end hooks. Would we still need these hooks if the summarization logic was removed from the trainer? Why doesn't this happen in on_train_start today? We don't have on_prevalidation_routine_start/end hooks: the necessity of these hooks for training isn't clear to me, and further deprecating these hooks could bring greater API clarity & simplification.
  https://github.com/PyTorchLightning/pytorch-lightning/blob/8a931732ae5135e3e55d9c7b7031d81837e5798a/pytorch_lightning/trainer/trainer.py#L1103-L1113

Pitch

A callback in Lightning naturally fits this extension purpose. It generalizes well across lightning modules, has great flexibility for when it can be called, and allows users to customize the summarization logic (e.g. integrate other libraries more easily).

• https://github.com/tyleryep/torchinfo

• https://github.com/facebookresearch/fvcore/blob/master/fvcore/nn/flop_count.py

• With this callback available, this logic can be removed from the core Trainer in order to be more pluggable:

• https://github.com/PyTorchLightning/pytorch-lightning/blob/6604fc1344e1b8a459c45a5a2157aa7fc60d950d/pytorch_lightning/trainer/trainer.py#L1000-L1004

Additional context

The model summary is by default enabled right now. This is likely the core issue we have to resolve as to whether this is opt-in or opt-out: #8478 (comment)

Seeking @edenafek @tchaton 's input on this

---

If you enjoy Lightning, check out our other projects! ⚡

• Metrics: Machine learning metrics for distributed, scalable PyTorch applications.

• Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, finetuning and solving problems with deep learning

• Bolts: Pretrained SOTA Deep Learning models, callbacks and more for research and production with PyTorch Lightning and PyTorch

• Lightning Transformers: Flexible interface for high performance research using SOTA Transformers leveraging Pytorch Lightning, Transformers, and Hydra.

[awaelchli]

@ananthsub the callback idea sounds nice but how do we make it so that a user can disable it? It's the same reason why we have the checkpoint_callback Trainer argument despite checkpointing being fully implemented as a callback. We set "good" defaults but we need a way to disable it too.

[ananthsub]

@ananthsub the callback idea sounds nice but how do we make it so that a user can disable it? It's the same reason why we have the checkpoint_callback Trainer argument despite checkpointing being fully implemented as a callback. We set "good" defaults but we need a way to disable it too.

This assumes the weights-summary is by default enabled. In this scenario, an incremental approach would be:

• Add the callback
• Update this connector to add the callback to the trainer list: https://github.com/PyTorchLightning/pytorch-lightning/blob/f3c5889aa3116bb8299a849f4f02dfb535b5624e/pytorch_lightning/trainer/connectors/callback_connector.py#L28-L65
• The trainer constructor arg remains the same. Users can already set weights_summary=None to disable it.

Or, we can make this opt-in. Remove the argument from the trainer constructor entirely, and enforce that users set this by instantiating a callback and passing it to the callbacks argument.

[awaelchli]

This assumes the weights-summary is by default enabled.

Yes this was my assumption and I hope we can keep that. I will definitely vote for that, but I am biased since I have been working on that summary, so I prefer if others could comment @PyTorchLightning/core-contributors . I strongly believe everyone working with ML models should be aware with how many parameters they are training with.

[ananthsub]

This assumes the weights-summary is by default enabled.

Yes this was my assumption and I hope we can keep that. I will definitely vote for that, but I am biased since I have been working on that summary, so I prefer if others could comment @PyTorchLightning/core-contributors . I strongly believe everyone working with ML models should be aware with how many parameters they are training with.

What this would likely also introduce is users needing to extend their callback from a particular summary base class in order for the trainer to validate whether it should add the default callback or not. This is the case for the Model
Checkpoint today.

However, the model checkpoint callback is not explicitly designed with extensibility in mind. See prior issues for offering a base interface:

This could limit what a custom summary could do in case the base class is too restrictive. I'd also like to avoid dependencies on inheritance wherever possible.

[ananthsub]

Would definitely like to hear @kaushikb11 's opinion given work on the Rich-based Model Summary

[kaushikb11]

@ananthsub I absolutely agree with you on introducing a new ModelSummary callback. We could easily extend the logic to a callback and make the Trainer cleaner. One more good point is that the user could extend the callback for different use cases. It could be extended for RichProgressBar as well, with updating the summarize functionality. As one can see in #9215, it feels hacky with the current design to extend it.

There are two parts to this issue as you mentioned. Regarding the second part proposal, I strongly believe we shouldn't deprecate the weights_summary argument. It's a good default to have as well as Users should be aware of the model parameters.

Also, with the current ModelSummary, we are able to support the below step as well

model = LitModel()
ModelSummary(model, max_depth=1)

Hence, we should be supporting both, the existing ModelSummary class and the new proposed callback.

[ananthsub]

@kaushikb11 @awaelchli - what do you think of this to mirror what was done for checkpoint_callback:

1. We create a new ModelSummary callback, which uses the summarize functionality recently moved to the utilities

2. We inject a ModelSummary callback into the callback constructor if none are set here and replace the hardcoded logic here. This way users could extend this callback for when the summary is generated, what information is included in the summary, or where the summary is outputted (e.g. stdout, file, etc).

3. On the Trainer, we type weights_summary as Union[str, bool] = True and mark the string values as deprecated, given the oncoming incompatibilities over mode vs max_depth. The expectation is that users who wish to customize the ModelSummary should pass in a customized ModelSummary callback through the callbacks argument.

[awaelchli]

I like that very much. This is the minimal functionality I wish we could keep.

Btw on a side note, I believe the best moving forward would be to have a model summary class with the only responsibility of collecting the summary data (like it is now) BUT not contain the logic for printing and visualization. I think it would be best if this would live in the callback. this way it will also be easier to customize things like rich logging etc while keeping the actual model summary untouched.

not sure if @kaushikb11 was already going in that direction, he might have

[ananthsub]

@awaelchli I fully agree regarding where the output of the summarization should go. Outputting the summary should live in the callback and not in the utils, as it does today