Skip to content

Fix ModelSummary.total_layer_params double-counting nested modules when max_depth != 1 - #22050

Open
deep1blue wants to merge 1 commit into
Lightning-AI:masterfrom
deep1blue:fix/model-summary-other-params-max-depth
Open

deep1blue wants to merge 1 commit into
Lightning-AI:masterfrom
deep1blue:fix/model-summary-other-params-max-depth

Conversation

@deep1blue

Copy link
Copy Markdown

What does this PR do?

When ModelSummary (or DeepSpeedSummary, which inherits total_layer_params) is initialized with max_depth > 1 or max_depth == -1, self._layer_summary contains both top-level child modules (e.g. net) and their nested submodules (e.g. net.0, net.1).

Because each LayerSummary.num_parameters calls self._module.parameters() recursively, a parent module's num_parameters already includes all parameters of its nested submodules. Summing self.param_nums across all entries in self._layer_summary in ModelSummary.total_layer_params therefore double/triple-counts nested submodule parameters whenever max_depth > 1 or max_depth == -1. As a consequence:

  1. summary.total_layer_params exceeds summary.total_parameters on any nested model with max_depth != 1.
  2. In _get_summary_data(), total_leftover_params = self.total_parameters - self.total_layer_params becomes <= 0, silently dropping the "other params" (LEFTOVER_PARAMS_NAME) row for direct nn.Parameter attributes on the LightningModule whenever max_depth > 1 or max_depth == -1.

Changes

  • Compute ModelSummary.total_layer_params by summing layer.num_parameters only across top-level child modules ("." not in name) in self._layer_summary.
  • Parameterize test_summary_data_with_non_layer_params in tests/tests_pytorch/utilities/test_model_summary.py over max_depth in [-1, 1, 2] with a nested submodule (nn.Sequential(nn.Linear(2, 2))) to verify that total_layer_params and the "other params" row remain accurate at all depths.

Acknowledgments

Credit to @awaelchli, @Borda, @carmocca, @SkafteNicki, and @GdoongMathew for their work on ModelSummary and DeepSpeedSummary.

…en max_depth != 1

Signed-off-by: Parth Vadera <parthvadera@gmail.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant